Lift Assignment: process lambda assignments in when correctly

So #KT-23346 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/branchedTransformations/BranchedFoldingUtils.kt

@@ -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) {

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

@@ -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" }
+    }
+}

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

@@ -0,0 +1,11 @@
+fun test(i: Int) {
+    var fn: () -> String
+
+    <caret>fn = if (i == 1) {
+        { "foo" }
+    } else if (i == 2) {
+        { "bar" }
+    } else {
+        { "baz" }
+    }
+}

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

@@ -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" }
+}

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

@@ -0,0 +1,11 @@
+fun test(i: Int) {
+    var fn: () -> String
+
+    <caret>fn = if (i == 1) {
+        { "foo" }
+    } else if (i == 2) {
+        { "bar" }
+    } else {
+        { "baz" }
+    }
+}

idea/testData/inspectionsLocal/liftOut/tryToAssignment/lambda.kt

@@ -0,0 +1,9 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>try {
+        fn = { "foo" }
+    } catch (e: Exception) {
+        fn = { "bar" }
+    }
+}

idea/testData/inspectionsLocal/liftOut/tryToAssignment/lambda.kt.after

@@ -0,0 +1,9 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>fn = try {
+        { "foo" }
+    } catch (e: Exception) {
+        { "bar" }
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToAssignment/lambda.kt

@@ -0,0 +1,12 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>when (b) {
+        true -> {
+            fn = { "foo" }
+        }
+        else -> {
+            fn = { "bar" }
+        }
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToAssignment/lambda.kt.after

@@ -0,0 +1,12 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>fn = when (b) {
+        true -> {
+            { "foo" }
+        }
+        else -> {
+            { "bar" }
+        }
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToAssignment/lambda2.kt

@@ -0,0 +1,8 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>when (b) {
+        true -> fn = { "foo" }
+        else -> fn = { "bar" }
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToAssignment/lambda2.kt.after

@@ -0,0 +1,12 @@
+fun test(b: Boolean) {
+    var fn: () -> String
+
+    <caret>fn = when (b) {
+        true -> {
+            { "foo" }
+        }
+        else -> {
+            { "bar" }
+        }
+    }
+}

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

@@ -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");