Lift Assignment: process lambda assignments in when correctly

So #KT-23346 Fixed
2018-03-24 16:47:40 +03:00
parent 7438d12ad7
commit 89981ecbbb
12 changed files with 138 additions and 2 deletions
@@ -162,12 +162,18 @@ object BranchedFoldingUtils {
    fun foldToAssignment(expression: KtExpression) {
        var lhs: KtExpression? = null
        var op: String? = null
+        val psiFactory = KtPsiFactory(expression)
        fun KtBinaryExpression.replaceWithRHS() {
            if (lhs == null || op == null) {
                lhs = left!!.copy() as KtExpression
                op = operationReference.text
            }
-            replace(right!!)
+            val rhs = right!!
+            if (rhs is KtLambdaExpression && this.parent !is KtBlockExpression) {
+                replace(psiFactory.createExpressionByPattern("{ $0 }", rhs))
+            } else {
+                replace(rhs)
+            }
        }
        fun lift(e: KtExpression?) {
            when (e) {
@@ -183,7 +189,7 @@ object BranchedFoldingUtils {
            }
        }
        lift(expression)
-        expression.replace(KtPsiFactory(expression).createExpressionByPattern("$0 $1 $2", lhs!!, op!!, expression))
+        expression.replace(psiFactory.createExpressionByPattern("$0 $1 $2", lhs!!, op!!, expression))
    }

    fun foldToReturn(expression: KtExpression) {
@@ -0,0 +1,11 @@
+fun test(i: Int) {
+    var fn: () -> String
+
+    <caret>if (i == 1) {
+        fn = { "foo" }
+    } else if (i == 2) {
+        fn = { "bar" }
+    } else {
+        fn = { "baz" }
+    }
+}
@@ -0,0 +1,11 @@
+fun test(i: Int) {
+    var fn: () -> String
+
+    <caret>fn = if (i == 1) {
+        { "foo" }
+    } else if (i == 2) {
+        { "bar" }
+    } else {
+        { "baz" }
+    }
+}
@@ -0,0 +1,10 @@
+fun test(i: Int) {
+    var fn: () -> String
+
+    <caret>if (i == 1)
+        fn = { "foo" }
+    else if (i == 2)
+        fn = { "bar" }
+    else
+        fn = { "baz" }
+}
@@ -0,0 +1,11 @@
+fun test(i: Int) {
+    var fn: () -> String
+
+    <caret>fn = if (i == 1) {
+        { "foo" }
+    } else if (i == 2) {
+        { "bar" }
+    } else {
+        { "baz" }
+    }
+}
@@ -0,0 +1,9 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>try {
+        fn = { "foo" }
+    } catch (e: Exception) {
+        fn = { "bar" }
+    }
+}
@@ -0,0 +1,9 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>fn = try {
+        { "foo" }
+    } catch (e: Exception) {
+        { "bar" }
+    }
+}
@@ -0,0 +1,12 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>when (b) {
+        true -> {
+            fn = { "foo" }
+        }
+        else -> {
+            fn = { "bar" }
+        }
+    }
+}
@@ -0,0 +1,12 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>fn = when (b) {
+        true -> {
+            { "foo" }
+        }
+        else -> {
+            { "bar" }
+        }
+    }
+}
@@ -0,0 +1,8 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>when (b) {
+        true -> fn = { "foo" }
+        else -> fn = { "bar" }
+    }
+}
@@ -0,0 +1,12 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>fn = when (b) {
+        true -> {
+            { "foo" }
+        }
+        else -> {
+            { "bar" }
+        }
+    }
+}
@@ -1795,6 +1795,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/innerIfTransformed.kt");
            }

+            @TestMetadata("lambda.kt")
+            public void testLambda() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/lambda.kt");
+            }
+
+            @TestMetadata("lambda2.kt")
+            public void testLambda2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/lambda2.kt");
+            }
+
            @TestMetadata("multipleAssignments.kt")
            public void testMultipleAssignments() throws Exception {
                runTest("idea/testData/inspectionsLocal/liftOut/ifToAssignment/multipleAssignments.kt");
@@ -1951,6 +1961,11 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                runTest("idea/testData/inspectionsLocal/liftOut/tryToAssignment/inner.kt");
            }

+            @TestMetadata("lambda.kt")
+            public void testLambda() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/tryToAssignment/lambda.kt");
+            }
+
            @TestMetadata("withUnmatchedAssignments.kt")
            public void testWithUnmatchedAssignments() throws Exception {
                runTest("idea/testData/inspectionsLocal/liftOut/tryToAssignment/withUnmatchedAssignments.kt");
@@ -2042,6 +2057,16 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                runTest("idea/testData/inspectionsLocal/liftOut/whenToAssignment/insideLoop.kt");
            }

+            @TestMetadata("lambda.kt")
+            public void testLambda() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/whenToAssignment/lambda.kt");
+            }
+
+            @TestMetadata("lambda2.kt")
+            public void testLambda2() throws Exception {
+                runTest("idea/testData/inspectionsLocal/liftOut/whenToAssignment/lambda2.kt");
+            }
+
            @TestMetadata("simpleWhen.kt")
            public void testSimpleWhen() throws Exception {
                runTest("idea/testData/inspectionsLocal/liftOut/whenToAssignment/simpleWhen.kt");