Extract Function: Fix call replacement for the case of lambda extraction

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractFunctionUtils.kt

@@ -71,6 +71,7 @@ import kotlin.properties.Delegates
 import org.jetbrains.jet.lang.cfg.pseudocodeTraverser.traverse
 import org.jetbrains.jet.lang.cfg.pseudocodeTraverser.TraversalOrder
 import org.jetbrains.jet.lang.resolve.bindingContextUtil.getTargetFunctionDescriptor
+import com.intellij.psi.PsiWhiteSpace
 
 private val DEFAULT_FUNCTION_NAME = "myFun"
 private val DEFAULT_RETURN_TYPE = KotlinBuiltIns.getInstance().getUnitType()
@@ -811,6 +812,27 @@ fun ExtractionDescriptor.generateFunction(
         }
     }
 
+    fun insertCall(anchor: PsiElement, wrappedCall: JetExpression) {
+        val firstExpression = extractionData.getExpressions().firstOrNull()
+        val enclosingCall = firstExpression?.getParent() as? JetCallExpression
+        if (enclosingCall == null || firstExpression !in enclosingCall.getFunctionLiteralArguments()) {
+            anchor.replace(wrappedCall)
+            return
+        }
+
+        val argumentListExt = JetPsiFactory.createCallArguments(project, "(${wrappedCall.getText()})")
+        val argumentList = enclosingCall.getValueArgumentList()
+        if (argumentList == null) {
+            (anchor.getPrevSibling() as? PsiWhiteSpace)?.let { it.delete() }
+            anchor.replace(argumentListExt)
+            return
+        }
+
+        val newArgText = (argumentList.getArguments() + argumentListExt.getArguments()).map { it.getText() }.joinToString(", ", "(", ")")
+        argumentList.replace(JetPsiFactory.createCallArguments(project, newArgText))
+        anchor.delete()
+    }
+
     fun makeCall(function: JetNamedFunction): JetNamedFunction {
         val anchor = extractionData.originalElements.first
         if (anchor == null) return function
@@ -849,7 +871,7 @@ fun ExtractionDescriptor.generateFunction(
             else ->
                 JetPsiFactory.createExpression(project, callText)
         }
-        anchor.replace(wrappedCall)
+        insertCall(anchor, wrappedCall)
 
         return function
     }

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaEmptyArgList.kt

@@ -0,0 +1,6 @@
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    t.check() <selection>{ it + 1 > 1 }</selection>
+}

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaEmptyArgList.kt.after

@@ -0,0 +1,10 @@
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    t.check(function())
+}
+
+fun function(): (Int) -> Boolean {
+    return { it + 1 > 1 }
+}

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNoArgList.kt

@@ -0,0 +1,6 @@
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    t.check <selection>{ it + 1 > 1 }</selection>
+}

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNoArgList.kt.after

@@ -0,0 +1,10 @@
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    t.check(function())
+}
+
+fun function(): (Int) -> Boolean {
+    return { it + 1 > 1 }
+}

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNonEmptyArgList.kt

@@ -0,0 +1,6 @@
+fun <T> Array<T>.check(a: Int, b: Int f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    t.check(1, 2) <selection>{ it + 1 > 1 }</selection>
+}

idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNonEmptyArgList.kt.after

@@ -0,0 +1,10 @@
+fun <T> Array<T>.check(a: Int, b: Int f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    t.check(1, 2, function())
+}
+
+fun function(): (Int) -> Boolean {
+    return { it + 1 > 1 }
+}

idea/tests/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/JetExtractionTestGenerated.java

@@ -447,6 +447,21 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                     doExtractFunctionTest("idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/simpleEvalExpr.kt");
                 }
                 
+                @TestMetadata("trailingLambdaEmptyArgList.kt")
+                public void testTrailingLambdaEmptyArgList() throws Exception {
+                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaEmptyArgList.kt");
+                }
+                
+                @TestMetadata("trailingLambdaNoArgList.kt")
+                public void testTrailingLambdaNoArgList() throws Exception {
+                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNoArgList.kt");
+                }
+                
+                @TestMetadata("trailingLambdaNonEmptyArgList.kt")
+                public void testTrailingLambdaNonEmptyArgList() throws Exception {
+                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/controlFlow/evaluateExpression/trailingLambdaNonEmptyArgList.kt");
+                }
+                
             }
             
             @TestMetadata("idea/testData/refactoring/extractFunction/controlFlow/outputValues")