From f73da4c0ea0e790503083e1cd2140696881da31e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christian=20Vonr=C3=BCti?= <christian@vonrueti.ch>
Date: Sat, 23 Apr 2022 11:04:48 +0200
Subject: [PATCH] Remove receiver when unwrapping trailing lambda KTIJ-21562

IntelliJ now transforms

    kotlin.run {
        println("abc")
    }

into

    println("abc")

when unwrapping.

closes https://github.com/JetBrains/intellij-community/pull/1996

GitOrigin-RevId: 7f547279382544eabe9252ec210857f62ba07815
---
 .../codeInsight/unwrap/KotlinLambdaUnwrapper.java     | 11 +++++++++--
 .../codeInsight/unwrap/UnwrapRemoveTestGenerated.java |  5 +++++
 .../unwrapLambda/lambdaWithQualifiedReceiver.kt       |  6 ++++++
 .../unwrapLambda/lambdaWithQualifiedReceiver.kt.after |  4 ++++
 4 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt
 create mode 100644 plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt.after

diff --git a/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/codeInsight/unwrap/KotlinLambdaUnwrapper.java b/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/codeInsight/unwrap/KotlinLambdaUnwrapper.java
index f2f72d4a4974..8261f6016be6 100644
--- a/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/codeInsight/unwrap/KotlinLambdaUnwrapper.java
+++ b/plugins/kotlin/idea/src/org/jetbrains/kotlin/idea/codeInsight/unwrap/KotlinLambdaUnwrapper.java
@@ -19,7 +19,14 @@ public class KotlinLambdaUnwrapper extends KotlinUnwrapRemoveBase {
         PsiElement parent = lambda.getParent();
 
         if (parent instanceof KtValueArgument) {
-            return PsiTreeUtil.getParentOfType(parent, KtCallExpression.class, true);
+            KtCallExpression callExpression = PsiTreeUtil.getParentOfType(parent, KtCallExpression.class, true);
+            if (callExpression != null) {
+                PsiElement parentParent = callExpression.getParent();
+                if (parentParent instanceof KtQualifiedExpression) {
+                    return (KtElement) parentParent;
+                }
+            }
+            return callExpression;
         }
 
         if (parent instanceof KtCallExpression) {
@@ -39,7 +46,7 @@ public class KotlinLambdaUnwrapper extends KotlinUnwrapRemoveBase {
 
         KtLambdaExpression lambda = (KtLambdaExpression) e;
         KtBlockExpression body = lambda.getBodyExpression();
-        KtElement enclosingElement = getLambdaEnclosingElement((KtLambdaExpression) e);
+        KtElement enclosingElement = getLambdaEnclosingElement(lambda);
 
         if (body == null || enclosingElement == null) return false;
 
diff --git a/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/codeInsight/unwrap/UnwrapRemoveTestGenerated.java b/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/codeInsight/unwrap/UnwrapRemoveTestGenerated.java
index 172c13e1393e..c49df19b0b40 100644
--- a/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/codeInsight/unwrap/UnwrapRemoveTestGenerated.java
+++ b/plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/codeInsight/unwrap/UnwrapRemoveTestGenerated.java
@@ -289,6 +289,11 @@ public abstract class UnwrapRemoveTestGenerated extends AbstractUnwrapRemoveTest
         public void testLambdaPropertyInBlock() throws Exception {
             runTest("testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaPropertyInBlock.kt");
         }
+
+        @TestMetadata("lambdaWithQualifiedReceiver.kt")
+        public void testLambdaWithQualifiedReceiver() throws Exception {
+            runTest("testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt");
+        }
     }
 
     @RunWith(JUnit3RunnerWithInners.class)
diff --git a/plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt b/plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt
new file mode 100644
index 000000000000..719493d97935
--- /dev/null
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt
@@ -0,0 +1,6 @@
+// OPTION: 0
+fun test() {
+    kotlin.run <caret>{
+        println(123)
+    }
+}
diff --git a/plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt.after b/plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt.after
new file mode 100644
index 000000000000..890c134851d2
--- /dev/null
+++ b/plugins/kotlin/idea/tests/testData/codeInsight/unwrapAndRemove/unwrapLambda/lambdaWithQualifiedReceiver.kt.after
@@ -0,0 +1,4 @@
+// OPTION: 0
+fun test() {
+    println(123)<caret>
+}