Amend "Terminate preceding call with semicolon" quickfix

Quickfix wasn't aware of cases, when expression to be fixed
 is a part of parent call expression or dot-qualified expression.
 Incorrect position for extracting trailing lambdas led to broken
 formatting.

 #KT-34694 Fixed

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

@@ -21,46 +21,127 @@ class AddSemicolonBeforeLambdaExpressionFix(element: KtLambdaExpression) : Kotli
     override fun getFamilyName(): String = text
 
     override fun invoke(project: Project, editor: Editor?, file: KtFile) {
-        val lambdaExpressionArgument = element?.parent?.safeAs<KtLambdaArgument>()
+        val lambdaExpressionArgument = element?.parent?.safeAs<KtLambdaArgument>() ?: return
-            ?: return
+        val callExpression = lambdaExpressionArgument.parent.safeAs<KtCallExpression>() ?: return
-        val callExpression = lambdaExpressionArgument.parent.safeAs<KtCallExpression>()
+
-            ?: return
+        val desiredEndOfCallExpression = lambdaExpressionArgument.findCorrectEndOfCall()
-        val desiredEndOfCallExpression =
+
-            PsiTreeUtil.findSiblingBackward(
-                lambdaExpressionArgument,
-                KtNodeTypes.LAMBDA_ARGUMENT,
-                null
-            ) ?: PsiTreeUtil.findSiblingBackward(
-                lambdaExpressionArgument,
-                KtNodeTypes.VALUE_ARGUMENT_LIST,
-                null
-            )
         desiredEndOfCallExpression?.let { endOfCall ->
-            makeNewExpressionsFromFollowingLambdas(callExpression, endOfCall)
+            val psiFactory = KtPsiFactory(project)
-            val semicolon = callExpression.parent.addAfter(
+
-                KtPsiFactory(project).createSemicolon(),
+            val addedSemicolon = when (val parent = callExpression.parent) {
-                callExpression
+                // Parent call is the call to the right, we want to give it this call's last lambda argument as a new receiver
-            )
+                is KtCallExpression -> liftTrailingNodesAndRelocateLastLambda(
-            editor?.caretModel?.moveToOffset(semicolon.startOffset)
+                    psiFactory, callExpression, endOfCall,
+                    lastLambdaAcceptor = parent,
+                    nodeBeforeSemicolon = callExpression
+                )
+                // Incorrect call is a part of dot-qualified expression before it
+                is KtDotQualifiedExpression -> {
+                    val grandparent = parent.parent
+                    if (grandparent.isCallOrDotExpression) {
+                        // Similar to call expression parent, but here correct lambda receiver is one level higher,
+                        // since our parent is the dot expression to the left
+                        liftTrailingNodesAndRelocateLastLambda(
+                            psiFactory, callExpression, endOfCall,
+                            lastLambdaAcceptor = grandparent,
+                            nodeBeforeSemicolon = parent
+                        )
+                    } else {
+                        // Extract trailing lambdas (and possible formatting) two levels higher,
+                        // before dot-qualified expression to the left of this call expression
+                        liftTrailingNodes(
+                            psiFactory, callExpression, endOfCall,
+                            addNodesAfter = parent
+                        )
+                    }
+                }
+                // Simple case: extract all trailing nodes right after call - it is a standalone call expression
+                else -> liftTrailingNodes(
+                    psiFactory, callExpression, endOfCall,
+                    addNodesAfter = callExpression
+                )
+            }
+            editor?.caretModel?.moveToOffset(addedSemicolon.startOffset)
         }
     }
 
-    private fun makeNewExpressionsFromFollowingLambdas(
+    private fun KtLambdaArgument.findCorrectEndOfCall() =
+        PsiTreeUtil.findSiblingBackward(this, KtNodeTypes.LAMBDA_ARGUMENT, null)
+            ?: PsiTreeUtil.findSiblingBackward(this, KtNodeTypes.VALUE_ARGUMENT_LIST, null)
+
+    private fun liftTrailingNodesAndRelocateLastLambda(
+        psiFactory: KtPsiFactory,
+        callExpression: KtCallExpression,
+        endOfCall: PsiElement,
+        lastLambdaAcceptor: PsiElement,
+        nodeBeforeSemicolon: PsiElement
+    ): PsiElement {
+        val (topCall, callHolder) = topLevelHolder(callExpression)
+        val semicolon = callHolder.addBefore(psiFactory.createSemicolon(), topCall)
+
+        makeNewExpressionsFromTrailingLambdas(callExpression, endOfCall, addNodesAfter = semicolon) { lastLambdaExpression ->
+            lastLambdaAcceptor.addAfter(lastLambdaExpression, nodeBeforeSemicolon)
+        }
+
+        callHolder.addBefore(nodeBeforeSemicolon, semicolon)
+        nodeBeforeSemicolon.delete()
+
+        return semicolon
+    }
+
+    private fun liftTrailingNodes(
+        psiFactory: KtPsiFactory,
+        callExpression: KtCallExpression,
+        endOfCall: PsiElement,
+        addNodesAfter: PsiElement
+    ): PsiElement {
+        makeNewExpressionsFromTrailingLambdas(callExpression, endOfCall, addNodesAfter)
+        return addNodesAfter.parent.addAfter(
+            psiFactory.createSemicolon(),
+            addNodesAfter
+        )
+    }
+
+    private val PsiElement.isCallOrDotExpression
+        get() = this is KtCallExpression || this is KtDotQualifiedExpression
+
+    data class TopExpressionAndHolder(val top: PsiElement, val holder: PsiElement)
+
+    private fun topLevelHolder(callExpression: KtCallExpression): TopExpressionAndHolder {
+        var me: PsiElement = callExpression
+        var parent: PsiElement = callExpression.parent
+        while (parent.isCallOrDotExpression) {
+            me = parent
+            parent = parent.parent
+        }
+        return TopExpressionAndHolder(me, parent)
+    }
+
+    private fun makeNewExpressionsFromTrailingLambdas(
         oldCallExpression: KtCallExpression,
-        endOfArguments: PsiElement
+        endOfArguments: PsiElement,
+        addNodesAfter: PsiElement,
+        lastLambdaHandler: ((PsiElement) -> Unit)? = null
     ) {
         var lastSibling = oldCallExpression.lastChild
-        val parentForCallExpression = oldCallExpression.parent
+        var lastLambdaWasProcessed = false
 
         while (lastSibling != endOfArguments) {
             when (lastSibling) {
-                is KtLambdaArgument -> parentForCallExpression.addAfter(
+                is KtLambdaArgument -> {
-                    lastSibling.getLambdaExpression() ?: lastSibling,
+                    val lambdaExpression: PsiElement = lastSibling.getLambdaExpression() ?: lastSibling
-                    oldCallExpression
+
-                )
+                    if (lastLambdaHandler != null && !lastLambdaWasProcessed) {
-                else -> parentForCallExpression.addAfter(
+                        lastLambdaWasProcessed = true
+                        lastLambdaHandler(lambdaExpression)
+                    } else {
+                        addNodesAfter.parent.addAfter(lambdaExpression, addNodesAfter)
+                    }
+                }
+                else -> addNodesAfter.parent.addAfter(
                     lastSibling,
-                    oldCallExpression
+                    addNodesAfter
                 )
             }
             lastSibling = lastSibling.prevSibling

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/dotExpressionAndCallExpression.kt

@@ -0,0 +1,10 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun doSomething() {}
+fun Any.foo() {}
+
+fun test() {
+    doSomething().foo()
+    // comment and formatting
+    { doSomething()<caret> }()
+}

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/dotExpressionAndCallExpression.kt.after

@@ -0,0 +1,10 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun doSomething() {}
+fun Any.foo() {}
+
+fun test() {
+    doSomething().foo();
+    // comment and formatting
+    { doSomething() }()
+}

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/multipleCalls.kt

@@ -0,0 +1,9 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun foo() {}
+
+fun test {
+    foo()
+    // comment and formatting
+    { { { foo() }<caret> } }()()()
+}

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/multipleCalls.kt.after

@@ -0,0 +1,9 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun foo() {}
+
+fun test {
+    foo();
+    // comment and formatting
+    { { { foo() } } }()()()
+}

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/onlyDotBefore.kt

@@ -0,0 +1,7 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun foo() {
+    15.toString()
+    // comment and formatting
+    {<caret>}
+}

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/onlyDotBefore.kt.after

@@ -0,0 +1,7 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun foo() {
+    15.toString();
+    // comment and formatting
+    {}
+}

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/twoLambdasTwoInvokes.kt

@@ -0,0 +1,7 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun foo() {
+    { "first" }.invoke()
+    // comment and formatting
+    {<caret> "second" }.invoke()
+}

idea/testData/quickfix/addSemicolonBeforeLambdaExpression/twoLambdasTwoInvokes.kt.after

@@ -0,0 +1,7 @@
+// "Terminate preceding call with semicolon" "true"
+
+fun foo() {
+    { "first" }.invoke();
+    // comment and formatting
+    { "second" }.invoke()
+}

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

@@ -1113,11 +1113,31 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/addSemicolonBeforeLambdaExpression/basic.kt");
         }
 
+        @TestMetadata("dotExpressionAndCallExpression.kt")
+        public void testDotExpressionAndCallExpression() throws Exception {
+            runTest("idea/testData/quickfix/addSemicolonBeforeLambdaExpression/dotExpressionAndCallExpression.kt");
+        }
+
+        @TestMetadata("multipleCalls.kt")
+        public void testMultipleCalls() throws Exception {
+            runTest("idea/testData/quickfix/addSemicolonBeforeLambdaExpression/multipleCalls.kt");
+        }
+
         @TestMetadata("multipleLambdas.kt")
         public void testMultipleLambdas() throws Exception {
             runTest("idea/testData/quickfix/addSemicolonBeforeLambdaExpression/multipleLambdas.kt");
         }
 
+        @TestMetadata("onlyDotBefore.kt")
+        public void testOnlyDotBefore() throws Exception {
+            runTest("idea/testData/quickfix/addSemicolonBeforeLambdaExpression/onlyDotBefore.kt");
+        }
+
+        @TestMetadata("twoLambdasTwoInvokes.kt")
+        public void testTwoLambdasTwoInvokes() throws Exception {
+            runTest("idea/testData/quickfix/addSemicolonBeforeLambdaExpression/twoLambdasTwoInvokes.kt");
+        }
+
         @TestMetadata("withComments.kt")
         public void testWithComments() throws Exception {
             runTest("idea/testData/quickfix/addSemicolonBeforeLambdaExpression/withComments.kt");