Change Signature: Fix processing of lambda arguments

compiler/frontend/src/org/jetbrains/kotlin/psi/JetTypeReference.java

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.psi;
 import com.intellij.lang.ASTNode;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.lexer.JetTokens;
 import org.jetbrains.kotlin.psi.stubs.KotlinPlaceHolderStub;
 import org.jetbrains.kotlin.psi.stubs.elements.JetStubElementTypes;
 
@@ -65,4 +66,8 @@ public class JetTypeReference extends JetElementImplStub<KotlinPlaceHolderStub<J
         }
         return answer != null ? answer : Collections.<JetAnnotationEntry>emptyList();
     }
+
+    public boolean hasParentheses() {
+        return findChildByType(JetTokens.LPAR) != null && findChildByType(JetTokens.LPAR) != null;
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/psi/TypeRefHelpers.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.psi.JetPsiFactory
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.PsiErrorElement
 import org.jetbrains.kotlin.psi.JetCallableDeclaration
+import org.jetbrains.kotlin.psi.JetFunctionType
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
 fun getTypeReference(declaration: JetCallableDeclaration): JetTypeReference? {
@@ -56,17 +57,25 @@ fun setTypeReference(declaration: JetCallableDeclaration, addAfter: PsiElement?,
 }
 
 fun JetCallableDeclaration.setReceiverTypeReference(typeRef: JetTypeReference?): JetTypeReference? {
+    val needParentheses = typeRef != null && typeRef.getTypeElement() is JetFunctionType && !typeRef.hasParentheses()
     val oldTypeRef = getReceiverTypeReference()
     if (typeRef != null) {
-        if (oldTypeRef != null) {
+        val newTypeRef =
-            return oldTypeRef.replace(typeRef) as JetTypeReference
+                if (oldTypeRef != null) {
-        }
+                    oldTypeRef.replace(typeRef) as JetTypeReference
-        else {
+                }
-            val anchor = getNameIdentifier()
+                else {
-            val newTypeRef = addBefore(typeRef, anchor) as JetTypeReference
+                    val anchor = getNameIdentifier()
-            addAfter(JetPsiFactory(getProject()).createDot(), newTypeRef)
+                    val newTypeRef = addBefore(typeRef, anchor) as JetTypeReference
-            return newTypeRef
+                    addAfter(JetPsiFactory(getProject()).createDot(), newTypeRef)
+                    newTypeRef
+                }
+        if (needParentheses) {
+            val argList = JetPsiFactory(getProject()).createCallArguments("()")
+            newTypeRef.addBefore(argList.getLeftParenthesis()!!, newTypeRef.getFirstChild())
+            newTypeRef.add(argList.getRightParenthesis()!!)
         }
+        return newTypeRef
     }
     else {
         if (oldTypeRef != null) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/util/callUtil.kt

@@ -76,7 +76,7 @@ public fun <C: ResolutionContext<C>> Call.hasUnresolvedArguments(context: Resolu
 
 public fun Call.getValueArgumentsInParentheses(): List<ValueArgument> = getValueArguments().filterArgsInParentheses()
 
-public fun JetCallExpression.getValueArgumentsInParentheses(): List<ValueArgument> = getValueArguments().filterArgsInParentheses()
+public fun JetCallElement.getValueArgumentsInParentheses(): List<ValueArgument> = getValueArguments().filterArgsInParentheses()
 
 public fun Call.getValueArgumentListOrElement(): JetElement = getValueArgumentList() ?: getCalleeExpression() ?: getCallElement()
 

idea/src/org/jetbrains/kotlin/idea/intentions/MoveLambdaOutsideParenthesesIntention.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import org.jetbrains.kotlin.idea.util.psiModificationUtil.moveLambdaOutsideParentheses
 import org.jetbrains.kotlin.psi.JetCallExpression
 import org.jetbrains.kotlin.psi.JetFunctionLiteralExpression
 import org.jetbrains.kotlin.psi.JetPsiFactory
@@ -37,19 +38,6 @@ public class MoveLambdaOutsideParenthesesIntention : JetSelfTargetingIntention<J
     }
 
     override fun applyTo(element: JetCallExpression, editor: Editor) {
-        val args = element.getValueArgumentsInParentheses()
+        element.moveLambdaOutsideParentheses()
-        val functionLiteral = args.last!!.getArgumentExpression()?.getText()
-        val calleeText = element.getCalleeExpression()?.getText()
-        if (calleeText == null || functionLiteral == null) return
-
-        val params = args.subList(0, args.size - 1).map { it.asElement().getText() ?: "" }.joinToString(", ", "(", ")")
-
-        val newCall =
-            if (params == "()") {
-                "$calleeText $functionLiteral"
-            } else {
-                "$calleeText$params $functionLiteral"
-            }
-        element.replace(JetPsiFactory(element).createExpression(newCall))
     }
 }

idea/src/org/jetbrains/kotlin/idea/refactoring/changeSignature/usages/JetFunctionCallUsage.java

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.refactoring.changeSignature.usages;
 
 import com.intellij.psi.PsiElement;
+import kotlin.KotlinPackage;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.CallableDescriptor;
@@ -25,6 +26,7 @@ import org.jetbrains.kotlin.descriptors.DeclarationDescriptor;
 import org.jetbrains.kotlin.idea.caches.resolve.ResolvePackage;
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.JetChangeInfo;
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.JetParameterInfo;
+import org.jetbrains.kotlin.idea.util.psiModificationUtil.PsiModificationUtilPackage;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilPackage;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
@@ -73,7 +75,7 @@ public class JetFunctionCallUsage extends JetUsageInfo<JetCallElement> {
     private void updateArgumentsAndReceiver(JetChangeInfo changeInfo, JetCallElement element) {
         JetValueArgumentList arguments = element.getValueArgumentList();
         assert arguments != null : "Argument list is expected: " + element.getText();
-        List<JetValueArgument> oldArguments = arguments.getArguments();
+        List<? extends ValueArgument> oldArguments = element.getValueArguments();
 
         boolean isNamedCall = oldArguments.size() > 1 && oldArguments.get(0).getArgumentName() != null;
         StringBuilder parametersBuilder = new StringBuilder("(");
@@ -99,7 +101,7 @@ public class JetFunctionCallUsage extends JetUsageInfo<JetCallElement> {
         parametersBuilder.append(')');
         JetValueArgumentList newArguments = JetPsiFactory(getProject()).createCallArguments(parametersBuilder.toString());
 
-        Map<Integer, JetValueArgument> argumentMap = getParamIndexToArgumentMap(changeInfo, oldArguments);
+        Map<Integer, ? extends ValueArgument> argumentMap = getParamIndexToArgumentMap(changeInfo, oldArguments);
         int argIndex = 0;
 
         JetParameterInfo newReceiverInfo = changeInfo.getReceiverParameterInfo();
@@ -131,37 +133,58 @@ public class JetFunctionCallUsage extends JetUsageInfo<JetCallElement> {
                 continue;
             }
 
-            JetValueArgument oldArgument = argumentMap.get(parameterInfo.getOldIndex());
+            ValueArgument oldArgument = argumentMap.get(parameterInfo.getOldIndex());
 
             if (oldArgument != null) {
                 JetValueArgumentName argumentName = oldArgument.getArgumentName();
                 JetSimpleNameExpression argumentNameExpression = argumentName != null ? argumentName.getReferenceExpression() : null;
                 changeArgumentName(argumentNameExpression, parameterInfo);
-                newArgument.replace(oldArgument);
+                //noinspection ConstantConditions
+                newArgument.replace(oldArgument instanceof JetFunctionLiteralArgument
+                                    ? psiFactory.createArgument(oldArgument.getArgumentExpression())
+                                    : oldArgument.asElement());
             }
             else if (parameterInfo.getDefaultValueForCall().isEmpty())
                 newArgument.delete();
         }
 
-        arguments.replace(newArguments);
+        List<JetFunctionLiteralArgument> lambdaArguments = element.getFunctionLiteralArguments();
-
+        boolean hasLambdaArgumentsBefore = !lambdaArguments.isEmpty();
-        if (newReceiverInfo == originalReceiverInfo) return;
+        if (hasLambdaArgumentsBefore) {
-
+            element.deleteChildRange(KotlinPackage.first(lambdaArguments), KotlinPackage.last(lambdaArguments));
-        PsiElement replacingElement = element;
-        if (newReceiverInfo != null) {
-            JetValueArgument receiverArgument = argumentMap.get(newReceiverInfo.getOldIndex());
-            JetExpression extensionReceiverExpression = receiverArgument != null ? receiverArgument.getArgumentExpression() : null;
-            String receiverText = extensionReceiverExpression != null
-                                  ? extensionReceiverExpression.getText()
-                                  : newReceiverInfo.getDefaultValueForCall();
-            if (receiverText.isEmpty()) {
-                receiverText = "_";
-            }
-
-            replacingElement = psiFactory.createExpression(receiverText + "." + element.getText());
         }
 
-        elementToReplace.replace(replacingElement);
+        JetValueArgument lastArgument = KotlinPackage.lastOrNull(newArguments.getArguments());
+        boolean hasTrailingLambdaInArgumentListAfter = lastArgument != null && lastArgument.getArgumentExpression() instanceof JetFunctionLiteralExpression;
+
+        arguments.replace(newArguments);
+
+        JetElement newElement = element;
+        if (newReceiverInfo != originalReceiverInfo) {
+            PsiElement replacingElement = element;
+            if (newReceiverInfo != null) {
+                ValueArgument receiverArgument = argumentMap.get(newReceiverInfo.getOldIndex());
+                JetExpression extensionReceiverExpression = receiverArgument != null ? receiverArgument.getArgumentExpression() : null;
+                String receiverText = extensionReceiverExpression != null
+                                      ? extensionReceiverExpression.getText()
+                                      : newReceiverInfo.getDefaultValueForCall();
+                if (receiverText.isEmpty()) {
+                    receiverText = "_";
+                }
+
+                replacingElement = psiFactory.createExpression(receiverText + "." + element.getText());
+            }
+
+            newElement = (JetElement) elementToReplace.replace(replacingElement);
+        }
+
+        if (hasLambdaArgumentsBefore && hasTrailingLambdaInArgumentListAfter) {
+            JetCallElement newCallElement =
+                    (JetCallElement) (newElement instanceof JetQualifiedExpression
+                                      ? ((JetQualifiedExpression) newElement).getSelectorExpression()
+                                      : newElement);
+            PsiModificationUtilPackage.moveLambdaOutsideParentheses(newCallElement);
+        }
     }
 
     @Nullable
@@ -177,11 +200,11 @@ public class JetFunctionCallUsage extends JetUsageInfo<JetCallElement> {
         return null;
     }
 
-    private static Map<Integer, JetValueArgument> getParamIndexToArgumentMap(JetChangeInfo changeInfo, List<JetValueArgument> oldArguments) {
+    private static Map<Integer, ? extends ValueArgument> getParamIndexToArgumentMap(JetChangeInfo changeInfo, List<? extends ValueArgument> oldArguments) {
-        Map<Integer, JetValueArgument> argumentMap = new HashMap<Integer, JetValueArgument>();
+        Map<Integer, ValueArgument> argumentMap = new HashMap<Integer, ValueArgument>();
 
         for (int i = 0; i < oldArguments.size(); i++) {
-            JetValueArgument argument = oldArguments.get(i);
+            ValueArgument argument = oldArguments.get(i);
             JetValueArgumentName argumentName = argument.getArgumentName();
             JetSimpleNameExpression argumentNameExpression = argumentName != null ? argumentName.getReferenceExpression() : null;
             String oldParameterName = argumentNameExpression != null ? argumentNameExpression.getReferencedName() : null;

idea/src/org/jetbrains/kotlin/idea/util/psiModificationUtil.kt

@@ -18,10 +18,7 @@ package org.jetbrains.kotlin.idea.util.psiModificationUtil
 
 import com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.JetCallExpression
+import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.JetExpression
-import org.jetbrains.kotlin.psi.JetFunctionLiteralArgument
-import org.jetbrains.kotlin.psi.JetPsiFactory
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.callUtil.getValueArgumentsInParentheses
@@ -77,3 +74,20 @@ fun JetFunctionLiteralArgument.moveInsideParenthesesAndReplaceWith(
     return oldCallExpression.replace(newCallExpression) as JetCallExpression
 }
 
+fun JetCallElement.moveLambdaOutsideParentheses() {
+    val args = getValueArgumentsInParentheses()
+    val functionLiteral = args.last!!.getArgumentExpression()?.getText()
+    val calleeText = getCalleeExpression()?.getText()
+    if (calleeText == null || functionLiteral == null) return
+
+    val params = args.subList(0, args.size - 1).map { it.asElement().getText() ?: "" }.joinToString(", ", "(", ")")
+
+    val newCall =
+            if (params == "()") {
+                "$calleeText $functionLiteral"
+            }
+            else {
+                "$calleeText$params $functionLiteral"
+            }
+    replace(JetPsiFactory(this).createExpression(newCall))
+}

idea/testData/refactoring/changeSignature/ConvertLambdaParameterToReceiverAfter.kt

@@ -0,0 +1,9 @@
+fun (() -> Int).foo(x: Int, y: Int): Int {
+    return x + this()
+}
+
+fun bar() {
+    {
+        3
+    }.foo(1, 2)
+}

idea/testData/refactoring/changeSignature/ConvertLambdaParameterToReceiverBefore.kt

@@ -0,0 +1,9 @@
+fun <caret>foo(x: Int, y: Int, cl: () -> Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1, 2) {
+        3
+    }
+}

idea/testData/refactoring/changeSignature/MoveLambdaParameterAfter.kt

@@ -0,0 +1,9 @@
+fun foo(x: Int, cl: () -> Int, y: Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1, {
+        3
+    }, 2)
+}

idea/testData/refactoring/changeSignature/MoveLambdaParameterBefore.kt

@@ -0,0 +1,9 @@
+fun <caret>foo(x: Int, y: Int, cl: () -> Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1, 2) {
+        3
+    }
+}

idea/testData/refactoring/changeSignature/RemoveLambdaParameterAfter.kt

@@ -0,0 +1,7 @@
+fun foo(x: Int, y: Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1, 2)
+}

idea/testData/refactoring/changeSignature/RemoveLambdaParameterBefore.kt

@@ -0,0 +1,9 @@
+fun <caret>foo(x: Int, y: Int, cl: () -> Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1, 2) {
+        3
+    }
+}

idea/testData/refactoring/changeSignature/RemoveLastNonLambdaParameterAfter.kt

@@ -0,0 +1,9 @@
+fun foo(cl: () -> Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo {
+        2
+    }
+}

idea/testData/refactoring/changeSignature/RemoveLastNonLambdaParameterBefore.kt

@@ -0,0 +1,9 @@
+fun <caret>foo(x: Int, cl: () -> Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1) {
+        2
+    }
+}

idea/testData/refactoring/changeSignature/RemoveParameterBeforeLambdaAfter.kt

@@ -0,0 +1,9 @@
+fun foo(x: Int, cl: () -> Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1) {
+        3
+    }
+}

idea/testData/refactoring/changeSignature/RemoveParameterBeforeLambdaBefore.kt

@@ -0,0 +1,9 @@
+fun <caret>foo(x: Int, y: Int, cl: () -> Int): Int {
+    return x + cl()
+}
+
+fun bar() {
+    foo(1, 2) {
+        3
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/changeSignature/JetChangeSignatureTest.java

@@ -800,6 +800,32 @@ public class JetChangeSignatureTest extends KotlinCodeInsightTestCase {
         doTest(changeInfo);
     }
 
+    public void testRemoveParameterBeforeLambda() throws Exception {
+        JetChangeInfo changeInfo = getChangeInfo();
+        changeInfo.removeParameter(1);
+        doTest(changeInfo);
+    }
+
+    public void testMoveLambdaParameter() throws Exception {
+        JetChangeInfo changeInfo = getChangeInfo();
+        JetParameterInfo[] newParameters = changeInfo.getNewParameters();
+        changeInfo.setNewParameter(1, newParameters[2]);
+        changeInfo.setNewParameter(2, newParameters[1]);
+        doTest(changeInfo);
+    }
+
+    public void testConvertLambdaParameterToReceiver() throws Exception {
+        JetChangeInfo changeInfo = getChangeInfo();
+        changeInfo.setReceiverParameterInfo(changeInfo.getNewParameters()[2]);
+        doTest(changeInfo);
+    }
+
+    public void testRemoveLambdaParameter() throws Exception {
+        JetChangeInfo changeInfo = getChangeInfo();
+        changeInfo.removeParameter(2);
+        doTest(changeInfo);
+    }
+
     @NotNull
     @Override
     protected String getTestDataPath() {