Convert to block body: when expression with Unit result is now handled as "return needed" #KT-12193 Fixed

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToBlockBodyIntention.kt

@@ -52,16 +52,25 @@ class ConvertToBlockBodyIntention : SelfTargetingIntention<KtDeclarationWithBody
     }
 
     companion object {
+        private fun KtExpression.whenAsResult(): Boolean = when (this) {
+            is KtWhenExpression -> true
+            is KtIfExpression -> (then?.whenAsResult() ?: false) || (`else`?.whenAsResult() ?: false)
+            is KtBlockExpression -> statements.lastOrNull()?.whenAsResult() ?: false
+            else -> false
+        }
+
         fun convert(declaration: KtDeclarationWithBody): KtDeclarationWithBody {
             val body = declaration.bodyExpression!!
 
+
             fun generateBody(returnsValue: Boolean): KtExpression {
                 val bodyType = body.analyze().getType(body)
+                val unitWhenAsResult = (bodyType == null || KotlinBuiltIns.isUnit(bodyType)) && body.whenAsResult()
                 val needReturn = returnsValue &&
                                  (bodyType == null || (!KotlinBuiltIns.isUnit(bodyType) && !KotlinBuiltIns.isNothing(bodyType)))
 
                 val factory = KtPsiFactory(declaration)
-                val statement = if (needReturn) factory.createExpressionByPattern("return $0", body) else body
+                val statement = if (needReturn || unitWhenAsResult) factory.createExpressionByPattern("return $0", body) else body
                 return factory.createSingleStatementBlock(statement)
             }
 

idea/testData/intentions/convertToBlockBody/ifWhenUnit.kt

@@ -0,0 +1,12 @@
+enum class AccessMode { READ, WRITE, EXECUTE }
+
+fun foo() {}
+
+fun whenExpr(access: AccessMode, arg: Boolean) = <caret>if (arg) {} else {
+    foo()
+    when (access) {
+        AccessMode.READ -> {}
+        AccessMode.WRITE -> {}
+        AccessMode.EXECUTE -> {}
+    }
+}

idea/testData/intentions/convertToBlockBody/ifWhenUnit.kt.after

@@ -0,0 +1,14 @@
+enum class AccessMode { READ, WRITE, EXECUTE }
+
+fun foo() {}
+
+fun whenExpr(access: AccessMode, arg: Boolean) {
+    return if (arg) {} else {
+        foo()
+        when (access) {
+            AccessMode.READ -> {}
+            AccessMode.WRITE -> {}
+            AccessMode.EXECUTE -> {}
+        }
+    }
+}

idea/testData/intentions/convertToBlockBody/whenUnit.kt

@@ -0,0 +1,7 @@
+enum class AccessMode { READ, WRITE, EXECUTE }
+
+fun whenExpr(access: AccessMode) = <caret>when (access) {
+    AccessMode.READ -> {}
+    AccessMode.WRITE -> {}
+    AccessMode.EXECUTE -> {}
+}

idea/testData/intentions/convertToBlockBody/whenUnit.kt.after

@@ -0,0 +1,9 @@
+enum class AccessMode { READ, WRITE, EXECUTE }
+
+fun whenExpr(access: AccessMode) {
+    return when (access) {
+        AccessMode.READ -> {}
+        AccessMode.WRITE -> {}
+        AccessMode.EXECUTE -> {}
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -4225,6 +4225,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("ifWhenUnit.kt")
+        public void testIfWhenUnit() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/ifWhenUnit.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("implicitlyNonUnitFun.kt")
         public void testImplicitlyNonUnitFun() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/implicitlyNonUnitFun.kt");
@@ -4284,6 +4290,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/valueIsAnonymousObject4.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("whenUnit.kt")
+        public void testWhenUnit() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToBlockBody/whenUnit.kt");
+            doTest(fileName);
+        }
     }
 
     @TestMetadata("idea/testData/intentions/convertToConcatenatedString")