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,63 +194,7 @@ public class JetFunctionDefinitionUsage<T extends PsiElement> extends JetUsageIn
         }
 
         if (changeInfo.isParameterSetOrOrderChanged()) {
-            int parametersCount = changeInfo.getNewParametersCount();
+            processParameterListWithStructuralChanges(changeInfo, element, parameterList, psiFactory);
-            boolean isLambda = element instanceof JetFunctionLiteral;
-
-            JetParameterList newParameterList = null;
-            if (isLambda) {
-                if (parametersCount == 0 && ((JetFunctionLiteral) element).getTypeReference() == null) {
-                    if (parameterList != null) {
-                        parameterList.delete();
-                        ASTNode arrowNode = ((JetFunctionLiteral)element).getArrowNode();
-                        if (arrowNode != null) {
-                            arrowNode.getPsi().delete();
-                        }
-                    }
-                }
-                else {
-                    newParameterList = psiFactory.createFunctionLiteralParameterList(changeInfo.getNewParametersSignature(
-                            (JetFunctionDefinitionUsage<PsiElement>) this, 0)
-                    );
-                }
-            }
-            else {
-                newParameterList = psiFactory.createParameterList(changeInfo.getNewParametersSignature(
-                        (JetFunctionDefinitionUsage<PsiElement>) this, 0)
-                );
-            }
-
-            if (newParameterList != null) {
-                if (parameterList != null) {
-                    newParameterList = (JetParameterList) parameterList.replace(newParameterList);
-                }
-                else {
-                    if (element instanceof JetClass) {
-                        PsiElement anchor = ((JetClass) element).getTypeParameterList();
-
-                        if (anchor == null) {
-                            anchor = ((JetClass) element).getNameIdentifier();
-                        }
-                        if (anchor != null) {
-                            newParameterList = (JetParameterList) element.addAfter(newParameterList, anchor);
-                        }
-                    }
-                    else if (isLambda) {
-                        //noinspection ConstantConditions
-                        JetFunctionLiteral functionLiteral = (JetFunctionLiteral) element;
-                        PsiElement anchor = functionLiteral.getLBrace();
-                        newParameterList = (JetParameterList) element.addAfter(newParameterList, anchor);
-                        if (functionLiteral.getArrowNode() == null) {
-                            Pair<PsiElement, PsiElement> whitespaceAndArrow = psiFactory.createWhitespaceAndArrow();
-                            element.addRangeAfter(whitespaceAndArrow.getFirst(), whitespaceAndArrow.getSecond(), newParameterList);
-                        }
-                    }
-                }
-            }
-
-            if (newParameterList != null) {
-                ShortenPackage.addToShorteningWaitSet(newParameterList);
-            }
         }
         else if (parameterList != null) {
             int paramIndex = 0;
@@ -265,6 +215,106 @@ public class JetFunctionDefinitionUsage<T extends PsiElement> extends JetUsageIn
         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) {
+            if (parametersCount == 0 && ((JetFunctionLiteral) element).getTypeReference() == null) {
+                if (parameterList != null) {
+                    parameterList.delete();
+                    ASTNode arrowNode = ((JetFunctionLiteral)element).getArrowNode();
+                    if (arrowNode != null) {
+                        arrowNode.getPsi().delete();
+                    }
+                    parameterList = null;
+                }
+            }
+            else {
+                newParameterList = psiFactory.createFunctionLiteralParameterList(changeInfo.getNewParametersSignature(
+                        (JetFunctionDefinitionUsage<PsiElement>) this, 0)
+                );
+                canReplaceEntireList = true;
+            }
+        }
+        else {
+            newParameterList = psiFactory.createParameterList(changeInfo.getNewParametersSignature(
+                    (JetFunctionDefinitionUsage<PsiElement>) this, 0)
+            );
+        }
+
+        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();
+
+                if (anchor == null) {
+                    anchor = ((JetClass) element).getNameIdentifier();
+                }
+                if (anchor != null) {
+                    newParameterList = (JetParameterList) element.addAfter(newParameterList, anchor);
+                }
+            }
+            else if (isLambda) {
+                //noinspection ConstantConditions
+                JetFunctionLiteral functionLiteral = (JetFunctionLiteral) element;
+                PsiElement anchor = functionLiteral.getLBrace();
+                newParameterList = (JetParameterList) element.addAfter(newParameterList, anchor);
+                if (functionLiteral.getArrowNode() == null) {
+                    Pair<PsiElement, PsiElement> whitespaceAndArrow = psiFactory.createWhitespaceAndArrow();
+                    element.addRangeAfter(whitespaceAndArrow.getFirst(), whitespaceAndArrow.getSecond(), newParameterList);
+                }
+            }
+        }
+
+        if (newParameterList != null) {
+            ShortenPackage.addToShorteningWaitSet(newParameterList);
+        }
+    }
+
+    private static JetParameterList replaceParameterListAndKeepDelimiters(JetParameterList parameterList, JetParameterList newParameterList) {
+        List<JetParameter> oldParameters = parameterList.getParameters();
+        List<JetParameter> newParameters = newParameterList.getParameters();
+        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)));
+        }
+
+        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));
+        }
+
+        return parameterList;
+    }
+
     private static void changeVisibility(JetChangeInfo changeInfo, PsiElement element) {
         JetModifierKeywordToken newVisibilityToken = JetRefactoringUtil.getVisibilityToken(changeInfo.getNewVisibility());
 

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() {