Support "Lift return out of when" for exhaustive when without else

So #KT-18852 Fixed

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

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.idea.intentions.branchedTransformations
 
+import org.jetbrains.kotlin.cfg.WhenChecker
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.intentions.branches
 import org.jetbrains.kotlin.lexer.KtTokens
@@ -60,7 +61,7 @@ object BranchedFoldingUtils {
         fun collectAssignmentsAndCheck(e: KtExpression?): Boolean = when (e) {
             is KtWhenExpression -> {
                 val entries = e.entries
-                KtPsiUtil.checkWhenExpressionHasSingleElse(e) &&
+                !e.hasMissingCases() &&
                 entries.isNotEmpty() &&
                 entries.all { entry ->
                     val assignment = getFoldableBranchedAssignment(entry.expression)?.run { assignments.add(this) }
@@ -128,7 +129,7 @@ object BranchedFoldingUtils {
         is KtWhenExpression -> {
             val entries = expression.entries
             when {
-                !KtPsiUtil.checkWhenExpressionHasSingleElse(expression) -> null
+                expression.hasMissingCases() -> null
                 entries.isEmpty() -> null
                 else -> getFoldableReturns(entries.map { it.expression })
             }
@@ -211,4 +212,7 @@ object BranchedFoldingUtils {
 
     private fun KtTryExpression.tryBlockAndCatchBodies(): List<KtExpression?> = listOf(tryBlock) + catchClauses.map { it.catchBody }
 
+    private fun KtWhenExpression.hasMissingCases(): Boolean =
+            !KtPsiUtil.checkWhenExpressionHasSingleElse(this) && WhenChecker.getMissingCases(this, this.analyze()).isNotEmpty()
+
 }

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

@@ -0,0 +1,16 @@
+// PROBLEM: none
+
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    var res: Int = 0
+
+    <caret>when (e) {
+        TestEnum.A -> res = 1
+        TestEnum.B -> res = 2
+    }
+
+    return res
+}

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

@@ -0,0 +1,15 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    var res: Int = 0
+
+    <caret>when (e) {
+        TestEnum.A -> res = 1
+        TestEnum.B -> res = 2
+        TestEnum.C -> res = 3
+    }
+
+    return res
+}

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

@@ -0,0 +1,15 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    var res: Int = 0
+
+    <caret>res = when (e) {
+        TestEnum.A -> 1
+        TestEnum.B -> 2
+        TestEnum.C -> 3
+    }
+
+    return res
+}

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

@@ -0,0 +1,15 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    var res: Int = 0
+
+    <caret>when (e) {
+        TestEnum.A -> res = 1
+        TestEnum.B -> res = 2
+        else -> res = 3
+    }
+
+    return res
+}

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

@@ -0,0 +1,15 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    var res: Int = 0
+
+    <caret>res = when (e) {
+        TestEnum.A -> 1
+        TestEnum.B -> 2
+        else -> 3
+    }
+
+    return res
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasMissingCase.kt

@@ -0,0 +1,13 @@
+// ERROR: A 'return' expression required in a function with a block body ('{...}')
+// PROBLEM: none
+
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    <caret>when (e) {
+        TestEnum.A -> return 1
+        TestEnum.B -> return 2
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCase.kt

@@ -0,0 +1,11 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    <caret>when (e) {
+        TestEnum.A -> return 1
+        TestEnum.B -> return 2
+        TestEnum.C -> return 3
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCase.kt.after

@@ -0,0 +1,11 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    <caret>return when (e) {
+        TestEnum.A -> 1
+        TestEnum.B -> 2
+        TestEnum.C -> 3
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCaseWithElse.kt

@@ -0,0 +1,11 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    <caret>when (e) {
+        TestEnum.A -> return 1
+        TestEnum.B -> return 2
+        else -> return 3
+    }
+}

idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCaseWithElse.kt.after

@@ -0,0 +1,11 @@
+enum class TestEnum{
+    A, B, C
+}
+
+fun test(e: TestEnum): Int {
+    <caret>return when (e) {
+        TestEnum.A -> 1
+        TestEnum.B -> 2
+        else -> 3
+    }
+}

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

@@ -873,6 +873,24 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToAssignment/simpleWhenWithoutTerminatingAssignment.kt");
                 doTest(fileName);
             }
+
+            @TestMetadata("whenHasMissingCase.kt")
+            public void testWhenHasMissingCase() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToAssignment/whenHasMissingCase.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("whenHasNoMissingCase.kt")
+            public void testWhenHasNoMissingCase() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToAssignment/whenHasNoMissingCase.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("whenHasNoMissingCaseWithElse.kt")
+            public void testWhenHasNoMissingCaseWithElse() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToAssignment/whenHasNoMissingCaseWithElse.kt");
+                doTest(fileName);
+            }
         }
 
         @TestMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn")
@@ -925,6 +943,24 @@ public class LocalInspectionTestGenerated extends AbstractLocalInspectionTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("whenHasMissingCase.kt")
+            public void testWhenHasMissingCase() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasMissingCase.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("whenHasNoMissingCase.kt")
+            public void testWhenHasNoMissingCase() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCase.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("whenHasNoMissingCaseWithElse.kt")
+            public void testWhenHasNoMissingCaseWithElse() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCaseWithElse.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("whenOneReturn.kt")
             public void testWhenOneReturn() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("idea/testData/inspectionsLocal/liftOut/whenToReturn/whenOneReturn.kt");