Change Signature: Retain formatting of existing parameters when parameter count/order changes

#KT-6412 Fixed

idea/src/org/jetbrains/jet/plugin/refactoring/changeSignature/usages/JetFunctionDefinitionUsage.java

@@ -17,8 +17,12 @@
 package org.jetbrains.jet.plugin.refactoring.changeSignature.usages;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiMethod;
+import com.intellij.psi.PsiWhiteSpace;
+import com.intellij.psi.util.PsiTreeUtil;
+import kotlin.KotlinPackage;
 import kotlin.Pair;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -41,6 +45,8 @@ import org.jetbrains.jet.plugin.refactoring.changeSignature.JetChangeInfo;
 import org.jetbrains.jet.plugin.refactoring.changeSignature.JetParameterInfo;
 import org.jetbrains.jet.plugin.refactoring.changeSignature.JetValVar;
 
+import java.util.List;
+
 import static org.jetbrains.jet.lang.psi.PsiPackage.JetPsiFactory;
 
 public class JetFunctionDefinitionUsage<T extends PsiElement> extends JetUsageInfo<T> {
@@ -188,8 +194,36 @@ public class JetFunctionDefinitionUsage<T extends PsiElement> extends JetUsageIn
         }
 
         if (changeInfo.isParameterSetOrOrderChanged()) {
+            processParameterListWithStructuralChanges(changeInfo, element, parameterList, psiFactory);
+        }
+        else if (parameterList != null) {
+            int paramIndex = 0;
+
+            for (JetParameter parameter : parameterList.getParameters()) {
+                JetParameterInfo parameterInfo = changeInfo.getNewParameters()[paramIndex];
+                changeParameter(paramIndex, parameter, parameterInfo);
+                paramIndex++;
+            }
+
+            ShortenPackage.addToShorteningWaitSet(parameterList);
+        }
+
+        if (changeInfo.isVisibilityChanged() && !JetPsiUtil.isLocal((JetDeclaration) element)) {
+            changeVisibility(changeInfo, element);
+        }
+
+        return true;
+    }
+
+    private void processParameterListWithStructuralChanges(
+            JetChangeInfo changeInfo,
+            PsiElement element,
+            JetParameterList parameterList,
+            JetPsiFactory psiFactory
+    ) {
         int parametersCount = changeInfo.getNewParametersCount();
         boolean isLambda = element instanceof JetFunctionLiteral;
+        boolean canReplaceEntireList = false;
 
         JetParameterList newParameterList = null;
         if (isLambda) {
@@ -200,12 +234,14 @@ public class JetFunctionDefinitionUsage<T extends PsiElement> extends JetUsageIn
                     if (arrowNode != null) {
                         arrowNode.getPsi().delete();
                     }
+                    parameterList = null;
                 }
             }
             else {
                 newParameterList = psiFactory.createFunctionLiteralParameterList(changeInfo.getNewParametersSignature(
                         (JetFunctionDefinitionUsage<PsiElement>) this, 0)
                 );
+                canReplaceEntireList = true;
             }
         }
         else {
@@ -214,10 +250,16 @@ public class JetFunctionDefinitionUsage<T extends PsiElement> extends JetUsageIn
             );
         }
 
-            if (newParameterList != null) {
+        if (newParameterList == null) return;
+
         if (parameterList != null) {
+            if (canReplaceEntireList) {
                 newParameterList = (JetParameterList) parameterList.replace(newParameterList);
             }
+            else {
+                newParameterList = replaceParameterListAndKeepDelimiters(parameterList, newParameterList);
+            }
+        }
         else {
             if (element instanceof JetClass) {
                 PsiElement anchor = ((JetClass) element).getTypeParameterList();
@@ -240,29 +282,37 @@ public class JetFunctionDefinitionUsage<T extends PsiElement> extends JetUsageIn
                 }
             }
         }
-            }
 
         if (newParameterList != null) {
             ShortenPackage.addToShorteningWaitSet(newParameterList);
         }
     }
-        else if (parameterList != null) {
-            int paramIndex = 0;
 
-            for (JetParameter parameter : parameterList.getParameters()) {
+    private static JetParameterList replaceParameterListAndKeepDelimiters(JetParameterList parameterList, JetParameterList newParameterList) {
-                JetParameterInfo parameterInfo = changeInfo.getNewParameters()[paramIndex];
+        List<JetParameter> oldParameters = parameterList.getParameters();
-                changeParameter(paramIndex, parameter, parameterInfo);
+        List<JetParameter> newParameters = newParameterList.getParameters();
-                paramIndex++;
+        int oldCount = oldParameters.size();
+        int newCount = newParameters.size();
+
+        int commonCount = Math.min(oldCount, newCount);
+        for (int i = 0; i < commonCount; i++) {
+            oldParameters.set(i, (JetParameter) oldParameters.get(i).replace(newParameters.get(i)));
         }
 
-            ShortenPackage.addToShorteningWaitSet(parameterList);
+        if (commonCount == 0) return (JetParameterList) parameterList.replace(newParameterList);
+
+        if (oldCount > commonCount) {
+            parameterList.deleteChildRange(oldParameters.get(commonCount - 1).getNextSibling(),
+                                           KotlinPackage.last(oldParameters));
+        }
+        else if (newCount > commonCount) {
+            parameterList.addRangeAfter(newParameters.get(commonCount - 1).getNextSibling(),
+                                        newParameterList.getLastChild().getPrevSibling(),
+                                        PsiTreeUtil.skipSiblingsBackward(parameterList.getLastChild(),
+                                                                         PsiWhiteSpace.class, PsiComment.class));
         }
 
-        if (changeInfo.isVisibilityChanged() && !JetPsiUtil.isLocal((JetDeclaration) element)) {
+        return parameterList;
-            changeVisibility(changeInfo, element);
-        }
-
-        return true;
     }
 
     private static void changeVisibility(JetChangeInfo changeInfo, PsiElement element) {

idea/testData/refactoring/changeSignature/ExpressionFunctionAfter.kt

@@ -1 +1 @@
-fun fun1(x1: Int, y1: Int): Int = x1 * 2 + fun1(x1, )
+fun fun1(x1: Int, y1: Int) : Int = x1 * 2 + fun1(x1, )

idea/testData/refactoring/changeSignature/NoDefaultValuesInOverridesAfter.kt

@@ -1,15 +1,18 @@
 trait T {
-    fun foo(a: Int = 1, b: String = "2")
+    fun foo(a: Int = 1,
+            b: String = "2")
 }
 
 open class A: T {
-    override fun foo(a: Int, b: String) {
+    override fun foo(a: Int,
+                     b: String) {
         throw UnsupportedOperationException()
     }
 }
 
 class B: A() {
-    override fun foo(a: Int, b: String) {
+    override fun foo(a: Int,
+                     b: String) {
         throw UnsupportedOperationException()
     }
 }

idea/testData/refactoring/changeSignature/ParameterListAddParamAfter.kt

@@ -0,0 +1,6 @@
+fun foo(
+        n: Int, s: String,
+        a: Any, l: Long
+) {
+
+}

idea/testData/refactoring/changeSignature/ParameterListAddParamBefore.kt

@@ -0,0 +1,6 @@
+fun <caret>foo(
+        n: Int, s: String,
+        a: Any
+) {
+
+}

idea/testData/refactoring/changeSignature/ParameterListRemoveAllParamsAfter.kt

@@ -0,0 +1,3 @@
+fun foo() {
+
+}

idea/testData/refactoring/changeSignature/ParameterListRemoveAllParamsBefore.kt

@@ -0,0 +1,6 @@
+fun <caret>foo(
+        n: Int, s: String,
+        a: Any
+) {
+
+}

idea/testData/refactoring/changeSignature/ParameterListRemoveParamAfter.kt

@@ -0,0 +1,5 @@
+fun <caret>foo(
+        n: Int, s: String
+) {
+
+}

idea/testData/refactoring/changeSignature/ParameterListRemoveParamBefore.kt

@@ -0,0 +1,6 @@
+fun <caret>foo(
+        n: Int, s: String,
+        a: Any
+) {
+
+}

idea/tests/org/jetbrains/jet/plugin/refactoring/changeSignature/JetChangeSignatureTest.java

@@ -318,8 +318,7 @@ public class JetChangeSignatureTest extends KotlinCodeInsightTestCase {
     public void testFunctionJavaUsagesAndOverridesAddParam() throws Exception {
         JetChangeInfo changeInfo = getChangeInfo();
         changeInfo.addParameter(new JetParameterInfo(-1, "s", KotlinBuiltIns.getInstance().getStringType(), null, "\"abc\"", null, null));
-        changeInfo.addParameter(new JetParameterInfo(-1, "o", KotlinBuiltIns.getInstance().getNullableAnyType(), null, "\"def\"", null,
+        changeInfo.addParameter(new JetParameterInfo(-1, "o", KotlinBuiltIns.getInstance().getNullableAnyType(), null, "\"def\"", null, null));
-                                                     null));
         doTest(changeInfo);
     }
 
@@ -563,6 +562,26 @@ public class JetChangeSignatureTest extends KotlinCodeInsightTestCase {
         doTest(changeInfo);
     }
 
+    public void testParameterListAddParam() throws Exception {
+        JetChangeInfo changeInfo = getChangeInfo();
+        changeInfo.addParameter(new JetParameterInfo(-1, "l", KotlinBuiltIns.getInstance().getLongType(), null, "", null, null));
+        doTest(changeInfo);
+    }
+
+    public void testParameterListRemoveParam() throws Exception {
+        JetChangeInfo changeInfo = getChangeInfo();
+        changeInfo.removeParameter(changeInfo.getNewParametersCount() - 1);
+        doTest(changeInfo);
+    }
+
+    public void testParameterListRemoveAllParams() throws Exception {
+        JetChangeInfo changeInfo = getChangeInfo();
+        for (int i = changeInfo.getNewParametersCount() - 1; i >= 0; i--) {
+            changeInfo.removeParameter(i);
+        }
+        doTest(changeInfo);
+    }
+
     @NotNull
     @Override
     protected String getTestDataPath() {