Convert to expression body: take nested whens into account

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

@@ -53,19 +53,13 @@ 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 || bodyType.isUnit()) && body.whenAsResult()
+                val unitWhenAsResult = (bodyType == null || bodyType.isUnit()) && body.resultingWhens().isNotEmpty()
                 val needReturn = returnsValue &&
                                  (bodyType == null || (!bodyType.isUnit() && !bodyType.isNothing()))
 

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

@@ -108,8 +108,11 @@ class ConvertToExpressionBodyIntention : SelfTargetingOffsetIndependentIntention
                 val expressionType = context.getType(statement) ?: return null
                 val isUnit = KotlinBuiltIns.isUnit(expressionType)
                 if (!isUnit && !KotlinBuiltIns.isNothing(expressionType)) return null
-                if (isUnit && statement is KtWhenExpression) {
-                    if (statement.elseExpression == null && context.get(BindingContext.EXHAUSTIVE_WHEN, statement) != true) return null
+                if (isUnit) {
+                    val resultingWhens = statement.resultingWhens()
+                    if (resultingWhens.any { it.elseExpression == null && context.get(BindingContext.EXHAUSTIVE_WHEN, it) != true}) {
+                        return null
+                    }
                 }
                 return statement
             }

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

@@ -131,6 +131,15 @@ fun KtExpression.negate(): KtExpression {
     return KtPsiFactory(this).createExpressionByPattern("!$0", this)
 }
 
+fun KtExpression.resultingWhens(): List<KtWhenExpression> = when (this) {
+    is KtWhenExpression -> listOf(this)
+    is KtIfExpression -> (then?.resultingWhens() ?: listOf()) + (`else`?.resultingWhens() ?: listOf())
+    is KtBinaryExpression -> (left?.resultingWhens() ?: listOf()) + (right?.resultingWhens() ?: listOf())
+    is KtUnaryExpression -> this.baseExpression?.resultingWhens() ?: listOf()
+    is KtBlockExpression -> statements.lastOrNull()?.resultingWhens() ?: listOf()
+    else -> listOf()
+}
+
 private fun KtExpression.specialNegation(): KtExpression? {
     val factory = KtPsiFactory(this)
     when (this) {

idea/testData/intentions/convertToExpressionBody/elvisWhenUnitNonExhaustive.kt

@@ -0,0 +1,9 @@
+// IS_APPLICABLE: false
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(access: AccessMode) {
+    <caret>println("result") ?: when (access) {
+        AccessMode.READ -> println("read")
+        AccessMode.WRITE -> println("write")
+    }
+}
+fun println(s: String) {}

idea/testData/intentions/convertToExpressionBody/ifWhenUnitExhaustive.kt

@@ -0,0 +1,19 @@
+// IS_APPLICABLE: true
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(mode: Boolean, access: AccessMode) {
+    <caret>if (mode) {
+        when (access) {
+            AccessMode.READ -> println("read")
+            AccessMode.WRITE -> println("write")
+            AccessMode.RW -> println("both")
+        }
+    }
+    else {
+        when (access) {
+            AccessMode.READ -> println("noread")
+            AccessMode.WRITE -> println("nowrite")
+            AccessMode.RW -> println("no both")
+        }
+    }
+}
+fun println(s: String) {}

idea/testData/intentions/convertToExpressionBody/ifWhenUnitExhaustive.kt.after

@@ -0,0 +1,17 @@
+// IS_APPLICABLE: true
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(mode: Boolean, access: AccessMode) = if (mode) {
+    when (access) {
+        AccessMode.READ -> println("read")
+        AccessMode.WRITE -> println("write")
+        AccessMode.RW -> println("both")
+    }
+}
+else {
+    when (access) {
+        AccessMode.READ -> println("noread")
+        AccessMode.WRITE -> println("nowrite")
+        AccessMode.RW -> println("no both")
+    }
+}
+fun println(s: String) {}

idea/testData/intentions/convertToExpressionBody/ifWhenUnitNonExhaustive.kt

@@ -0,0 +1,18 @@
+// IS_APPLICABLE: false
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(mode: Boolean, access: AccessMode) {
+    <caret>if (mode) {
+        when (access) {
+            AccessMode.READ -> println("read")
+            AccessMode.WRITE -> println("write")
+        }
+    }
+    else {
+        when (access) {
+            AccessMode.READ -> println("noread")
+            AccessMode.WRITE -> println("nowrite")
+            AccessMode.RW -> println("no both")
+        }
+    }
+}
+fun println(s: String) {}

idea/testData/intentions/convertToExpressionBody/lambdaWhenUnitNonExhaustive.kt

@@ -0,0 +1,13 @@
+// IS_APPLICABLE: true
+enum class AccessMode { READ, WRITE, RW }
+fun <T> run(f: () -> T) = f()
+fun whenExpr(access: AccessMode) {
+    <caret>run {
+        println("run")
+        when (access) {
+            AccessMode.READ -> println("read")
+            AccessMode.WRITE -> println("write")
+        }
+    }
+}
+fun println(s: String) {}

idea/testData/intentions/convertToExpressionBody/lambdaWhenUnitNonExhaustive.kt.after

@@ -0,0 +1,11 @@
+// IS_APPLICABLE: true
+enum class AccessMode { READ, WRITE, RW }
+fun <T> run(f: () -> T) = f()
+fun whenExpr(access: AccessMode) = run {
+    println("run")
+    when (access) {
+        AccessMode.READ -> println("read")
+        AccessMode.WRITE -> println("write")
+    }
+}
+fun println(s: String) {}

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

@@ -4489,6 +4489,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("elvisWhenUnitNonExhaustive.kt")
+        public void testElvisWhenUnitNonExhaustive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/elvisWhenUnitNonExhaustive.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("emptyList.kt")
         public void testEmptyList() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/emptyList.kt");
@@ -4561,6 +4567,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("ifWhenUnitExhaustive.kt")
+        public void testIfWhenUnitExhaustive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/ifWhenUnitExhaustive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("ifWhenUnitNonExhaustive.kt")
+        public void testIfWhenUnitNonExhaustive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/ifWhenUnitNonExhaustive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("lambdaWhenUnitNonExhaustive.kt")
+        public void testLambdaWhenUnitNonExhaustive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/lambdaWhenUnitNonExhaustive.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("multipleStatements.kt")
         public void testMultipleStatements() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/multipleStatements.kt");