Replace with dot call quickfix: don't break formatting

#KT-25697 Fixed

idea/src/org/jetbrains/kotlin/idea/quickfix/ReplaceCallFix.kt

@@ -19,6 +19,8 @@ package org.jetbrains.kotlin.idea.quickfix
 import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.psi.*
@@ -45,15 +47,25 @@ abstract class ReplaceCallFix(
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
         val element = element ?: return
         val elvis = elvisOrEmpty(notNullNeeded)
+        val betweenReceiverAndOperation = element.elementsBetweenReceiverAndOperation().joinToString(separator = "") { it.text }
         val newExpression = KtPsiFactory(element).createExpressionByPattern(
-            "$0$operation$1$elvis",
+            "$0$betweenReceiverAndOperation$operation$1$elvis",
-            element.receiverExpression, element.selectorExpression!!
+            element.receiverExpression,
+            element.selectorExpression!!
         )
         val replacement = element.replace(newExpression)
         if (notNullNeeded) {
             replacement.moveCaretToEnd(editor, project)
         }
     }
+
+    private fun KtQualifiedExpression.elementsBetweenReceiverAndOperation(): List<PsiElement> {
+        val receiver = receiverExpression
+        val operation = operationTokenNode as? PsiElement ?: return emptyList()
+        val start = receiver.nextSibling?.takeIf { it != operation } ?: return emptyList()
+        val end = operation.prevSibling?.takeIf { it != receiver } ?: return emptyList()
+        return PsiTreeUtil.getElementsOfRange(start, end)
+    }
 }
 
 class ReplaceImplicitReceiverCallFix(

idea/testData/quickfix/replaceWithDotCall/comment.kt

@@ -0,0 +1,7 @@
+// "Replace with dot call" "true"
+// WITH_RUNTIME
+fun foo(a: String) {
+    val b = a // comment1
+            // comment2
+            ?.<caret>length
+}

idea/testData/quickfix/replaceWithDotCall/comment.kt.after

@@ -0,0 +1,7 @@
+// "Replace with dot call" "true"
+// WITH_RUNTIME
+fun foo(a: String) {
+    val b = a // comment1
+            // comment2
+            .length
+}

idea/testData/quickfix/replaceWithDotCall/lineBreak.kt

@@ -0,0 +1,6 @@
+// "Replace with dot call" "true"
+// WITH_RUNTIME
+fun foo(a: String) {
+    val b = a
+            ?.<caret>length
+}

idea/testData/quickfix/replaceWithDotCall/lineBreak.kt.after

@@ -0,0 +1,6 @@
+// "Replace with dot call" "true"
+// WITH_RUNTIME
+fun foo(a: String) {
+    val b = a
+            .length
+}

idea/testData/quickfix/replaceWithSafeCall/comment.kt

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a // comment1
+            // comment2
+            .<caret>length
+}

idea/testData/quickfix/replaceWithSafeCall/comment.kt.after

@@ -0,0 +1,7 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a // comment1
+            // comment2
+            ?.length
+}

idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a
+            .<caret>length
+}

idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt.after

@@ -0,0 +1,6 @@
+// "Replace with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a
+            ?.length
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/comment.kt

@@ -0,0 +1,9 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a // comment1
+            // comment2
+            .let {
+                it<caret>.length
+            }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/comment.kt.after

@@ -0,0 +1,9 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a // comment1
+            // comment2
+            ?.let {
+                it.length
+            }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/lineBreak.kt

@@ -0,0 +1,8 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a
+            .let {
+                it<caret>.length
+            }
+}

idea/testData/quickfix/replaceWithSafeCallForScopeFunction/lineBreak.kt.after

@@ -0,0 +1,8 @@
+// "Replace scope function with safe (?.) call" "true"
+// WITH_RUNTIME
+fun foo(a: String?) {
+    val b = a
+            ?.let {
+                it.length
+            }
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -10883,11 +10883,21 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/quickfix/replaceWithDotCall"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
         }
 
+        @TestMetadata("comment.kt")
+        public void testComment() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithDotCall/comment.kt");
+        }
+
         @TestMetadata("functionCall.kt")
         public void testFunctionCall() throws Exception {
             runTest("idea/testData/quickfix/replaceWithDotCall/functionCall.kt");
         }
 
+        @TestMetadata("lineBreak.kt")
+        public void testLineBreak() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithDotCall/lineBreak.kt");
+        }
+
         @TestMetadata("normal.kt")
         public void testNormal() throws Exception {
             runTest("idea/testData/quickfix/replaceWithDotCall/normal.kt");
@@ -10941,6 +10951,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/replaceWithSafeCall/assignmentToProperty.kt");
         }
 
+        @TestMetadata("comment.kt")
+        public void testComment() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/comment.kt");
+        }
+
         @TestMetadata("expression.kt")
         public void testExpression() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCall/expression.kt");
@@ -10971,6 +10986,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/replaceWithSafeCall/letWithParameter.kt");
         }
 
+        @TestMetadata("lineBreak.kt")
+        public void testLineBreak() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCall/lineBreak.kt");
+        }
+
         @TestMetadata("noReplaceWithSafeCallForImplicitReceiver.kt")
         public void testNoReplaceWithSafeCallForImplicitReceiver() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCall/noReplaceWithSafeCallForImplicitReceiver.kt");
@@ -11039,6 +11059,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/assignment.kt");
         }
 
+        @TestMetadata("comment.kt")
+        public void testComment() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/comment.kt");
+        }
+
         @TestMetadata("let.kt")
         public void testLet() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/let.kt");
@@ -11059,6 +11084,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/letWithWrongParam.kt");
         }
 
+        @TestMetadata("lineBreak.kt")
+        public void testLineBreak() throws Exception {
+            runTest("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/lineBreak.kt");
+        }
+
         @TestMetadata("notInsideScope.kt")
         public void testNotInsideScope() throws Exception {
             runTest("idea/testData/quickfix/replaceWithSafeCallForScopeFunction/notInsideScope.kt");