Lift assignment out: if last statement is lambda, enclose it in parentheses if necessary

#KT-38155 Fixed

compiler/psi/src/org/jetbrains/kotlin/psi/KtPsiUtil.java

@@ -450,7 +450,7 @@ public class KtPsiUtil {
 
         if (innerExpression instanceof KtLambdaExpression) {
             PsiElement prevSibling = PsiTreeUtil.skipWhitespacesAndCommentsBackward(currentInner);
-            if (prevSibling != null && prevSibling.getText().endsWith(KtTokens.RPAR.getValue())) return true;
+            if (endWithParenthesisOrCallExpression(prevSibling)) return true;
         }
 
         if (parentElement instanceof KtCallExpression && currentInner == ((KtCallExpression) parentElement).getCalleeExpression()) {
@@ -557,6 +557,15 @@ public class KtPsiUtil {
         return innerPriority < parentPriority;
     }
 
+    private static boolean endWithParenthesisOrCallExpression(PsiElement element) {
+        if (element == null) return false;
+        if (element.getText().endsWith(KtTokens.RPAR.getValue()) || element instanceof KtCallExpression) return true;
+        PsiElement[] children = element.getChildren();
+        int length = children.length;
+        if (length == 0) return false;
+        return endWithParenthesisOrCallExpression(children[length - 1]);
+    }
+
     private static boolean isKeepBinaryExpressionParenthesized(KtBinaryExpression expression) {
         PsiElement expr = expression.getFirstChild();
         while (expr != null) {

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace.kt

@@ -0,0 +1,12 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    <caret>if (i == 1) {
+        f = { true }
+    } else {
+        foo { i }
+        f = { false }
+    }
+    f()
+}
+
+fun foo(f: () -> Int) {}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace.kt.after

@@ -0,0 +1,12 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    f = if (i == 1) {
+        { true }
+    } else {
+        foo { i }
+        ({ false })
+    }
+    f()
+}
+
+fun foo(f: () -> Int) {}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace2.kt

@@ -0,0 +1,14 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    <caret>if (i == 1) {
+        f = { true }
+    } else {
+        val foo = Foo().foo { i } // comment
+        f = { false }
+    }
+    f()
+}
+
+class Foo {
+    fun foo(f: () -> Int) {}
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace2.kt.after

@@ -0,0 +1,14 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    f = if (i == 1) {
+        { true }
+    } else {
+        val foo = Foo().foo { i } // comment
+        ({ false })
+    }
+    f()
+}
+
+class Foo {
+    fun foo(f: () -> Int) {}
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace3.kt

@@ -0,0 +1,10 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    <caret>if (i == 1) {
+        f = { true }
+    } else {
+        val g: () -> Boolean = { false }
+        f = { g() }
+    }
+    f()
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace3.kt.after

@@ -0,0 +1,10 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    f = if (i == 1) {
+        { true }
+    } else {
+        val g: () -> Boolean = { false }
+        { g() }
+    }
+    f()
+}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightParenthesis.kt

@@ -0,0 +1,12 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    <caret>if (i == 1) {
+        f = { true }
+    } else {
+        val foo = foo()
+        f = { false }
+    }
+    f()
+}
+
+fun foo() {}

idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightParenthesis.kt.after

@@ -0,0 +1,12 @@
+fun test(i: Int) {
+    val f: () -> Boolean
+    f = if (i == 1) {
+        { true }
+    } else {
+        val foo = foo()
+        ({ false })
+    }
+    f()
+}
+
+fun foo() {}

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

@@ -5603,6 +5603,26 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
             }
 
+            @TestMetadata("afterRightBrace.kt")
+            public void testAfterRightBrace() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace.kt");
+            }
+
+            @TestMetadata("afterRightBrace2.kt")
+            public void testAfterRightBrace2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace2.kt");
+            }
+
+            @TestMetadata("afterRightBrace3.kt")
+            public void testAfterRightBrace3() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightBrace3.kt");
+            }
+
+            @TestMetadata("afterRightParenthesis.kt")
+            public void testAfterRightParenthesis() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/afterRightParenthesis.kt");
+            }
+
             public void testAllFilesPresentInIfToAssignment() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/inspectionsLocal/liftOut/ifToAssignment"), Pattern.compile("^([\\w\\-_]+)\\.(kt|kts)$"), null, true);
             }