SimplifyCallChainFix: Keep comments in place #KT-22552 Fixed

idea/src/org/jetbrains/kotlin/idea/inspections/collections/SimplifyCallChainFix.kt

@@ -21,8 +21,8 @@ import com.intellij.codeInspection.ProblemDescriptor
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.idea.core.ShortenReferences
 import org.jetbrains.kotlin.idea.core.replaced
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
 
 class SimplifyCallChainFix(private val newCallText: String) : LocalQuickFix {
     private val shortenedText = newCallText.split("(").joinToString(separator = "(") {
@@ -49,12 +49,29 @@ class SimplifyCallChainFix(private val newCallText: String) : LocalQuickFix {
             val firstCallExpression = AbstractCallChainChecker.getCallExpression(firstExpression) ?: return
             val secondCallExpression = secondQualifiedExpression.selectorExpression as? KtCallExpression ?: return
 
-            val lastArgumentPrefix = if (newCallText.startsWith("joinTo")) "transform = " else ""
+            val lastArgumentName = if (newCallText.startsWith("joinTo")) Name.identifier("transform") else null
-            val arguments = secondCallExpression.valueArgumentList?.arguments.orEmpty().map { it.text } +
+            if (lastArgumentName != null) {
-                    firstCallExpression.valueArgumentList?.arguments.orEmpty().map { "$lastArgumentPrefix${it.text}" }
+                val lastArgument = firstCallExpression.valueArgumentList?.arguments?.singleOrNull()
+                val argumentExpression = lastArgument?.getArgumentExpression()
+                if (argumentExpression != null) {
+                    lastArgument.replace(factory.createArgument(argumentExpression, lastArgumentName))
+                }
+            }
+            val firstCallArgumentList = firstCallExpression.valueArgumentList
+            val firstCallArguments = firstCallArgumentList?.arguments
+            val secondCallArgumentList = secondCallExpression.valueArgumentList
+            val secondCallArguments = secondCallArgumentList?.arguments
+
+            val argumentsText = when {
+                secondCallArguments?.isNotEmpty() == true && firstCallArguments?.isNotEmpty() == true -> {
+                    "${secondCallArgumentList.text.removeSuffix(")")}, ${firstCallArgumentList.text.removePrefix("(")}"
+                }
+                secondCallArguments?.isNotEmpty() == true -> secondCallArgumentList.text
+                firstCallArguments?.isNotEmpty() == true -> firstCallArgumentList.text
+                else -> ""
+            }
             val lambdaExpression = firstCallExpression.lambdaArguments.singleOrNull()?.getLambdaExpression()
 
-            val argumentsText = arguments.ifNotEmpty { joinToString(prefix = "(", postfix = ")") } ?: ""
             val newQualifiedExpression = if (lambdaExpression != null) factory.createExpressionByPattern(
                 "$0$1$2 $3 $4",
                 receiverExpression ?: "",

idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToStringWithReferenceAndComment.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+
+val x = listOf(1, 2, 3).<caret>map(
+        // comment1
+        Int::toString
+).joinToString(
+        // comment2
+        prefix = "= ",
+        // comment3
+        separator = " + "
+        // comment4
+)

idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToStringWithReferenceAndComment.kt.after

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+
+val x = listOf(1, 2, 3).joinToString(
+        // comment2
+        prefix = "= ",
+        // comment3
+        separator = " + "
+        // comment4
+        ,
+        // comment1
+        transform = Int::toString
+)

idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToWithComment.kt

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+
+val sb = StringBuilder()
+val x = listOf(1, 2, 3).<caret>map { "$it*$it" }.joinTo(
+        // comment1
+        buffer = sb,
+        // comment2
+        prefix = "= ",
+        // comment3
+        separator = " + "
+        // comment4
+)

idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToWithComment.kt.after

@@ -0,0 +1,12 @@
+// WITH_RUNTIME
+
+val sb = StringBuilder()
+val x = listOf(1, 2, 3).joinTo(
+        // comment1
+        buffer = sb,
+        // comment2
+        prefix = "= ",
+        // comment3
+        separator = " + "
+        // comment4
+) { "$it*$it" }

idea/testData/inspectionsLocal/collections/simplifiableCallChain/listOfNotNullWithComment.kt

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+
+fun test() {
+    <caret>listOf(
+            true, // comment1
+            null // comment2
+    ).filterNotNull().first()
+}

idea/testData/inspectionsLocal/collections/simplifiableCallChain/listOfNotNullWithComment.kt.after

@@ -0,0 +1,8 @@
+// WITH_RUNTIME
+
+fun test() {
+    <caret>listOfNotNull(
+            true, // comment1
+            null // comment2
+    ).first()
+}

idea/tests/org/jetbrains/kotlin/idea/inspections/LocalInspectionTestGenerated.java

@@ -1113,16 +1113,31 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 runTest("idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToStringWithReference.kt");
             }
 
+            @TestMetadata("joinToStringWithReferenceAndComment.kt")
+            public void testJoinToStringWithReferenceAndComment() throws Exception {
+                runTest("idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToStringWithReferenceAndComment.kt");
+            }
+
             @TestMetadata("joinToStringWithReferenceFake.kt")
             public void testJoinToStringWithReferenceFake() throws Exception {
                 runTest("idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToStringWithReferenceFake.kt");
             }
 
+            @TestMetadata("joinToWithComment.kt")
+            public void testJoinToWithComment() throws Exception {
+                runTest("idea/testData/inspectionsLocal/collections/simplifiableCallChain/joinToWithComment.kt");
+            }
+
             @TestMetadata("listOfNotNull.kt")
             public void testListOfNotNull() throws Exception {
                 runTest("idea/testData/inspectionsLocal/collections/simplifiableCallChain/listOfNotNull.kt");
             }
 
+            @TestMetadata("listOfNotNullWithComment.kt")
+            public void testListOfNotNullWithComment() throws Exception {
+                runTest("idea/testData/inspectionsLocal/collections/simplifiableCallChain/listOfNotNullWithComment.kt");
+            }
+
             @TestMetadata("mapNotNull.kt")
             public void testMapNotNull() throws Exception {
                 runTest("idea/testData/inspectionsLocal/collections/simplifiableCallChain/mapNotNull.kt");