From 2350caf1771c4f30a0c704284a78976c663e602d Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <Mikhail.Glukhikh@jetbrains.com>
Date: Fri, 27 May 2016 13:33:14 +0300
Subject: [PATCH] Convert to expression body is forbidden on single
 non-exhaustive when statement with Unit result #KT-12502 Fixed

---
 .../ConvertToExpressionBodyIntention.kt        | 11 +++++++++--
 .../whenUnitExhaustive.kt                      |  9 +++++++++
 .../whenUnitExhaustive.kt.after                |  7 +++++++
 .../whenUnitNonExhaustive.kt                   |  9 +++++++++
 .../whenUnitWithElse.kt                        |  9 +++++++++
 .../whenUnitWithElse.kt.after                  |  7 +++++++
 .../intentions/IntentionTestGenerated.java     | 18 ++++++++++++++++++
 7 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt
 create mode 100644 idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt.after
 create mode 100644 idea/testData/intentions/convertToExpressionBody/whenUnitNonExhaustive.kt
 create mode 100644 idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt
 create mode 100644 idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToExpressionBodyIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToExpressionBodyIntention.kt
index 2c71e4c2d10..303c11d8c90 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToExpressionBodyIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertToExpressionBodyIntention.kt
@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.utils.addToStdlib.check
 
 class ConvertToExpressionBodyIntention : SelfTargetingOffsetIndependentIntention<KtDeclarationWithBody>(
@@ -102,8 +103,14 @@ class ConvertToExpressionBodyIntention : SelfTargetingOffsetIndependentIntention
 
             else  -> {
                 if (statement is KtBinaryExpression && statement.operationToken in KtTokens.ALL_ASSIGNMENTS) return null // assignment does not have value
-                val expressionType = statement.analyze().getType(statement) ?: return null
-                if (!KotlinBuiltIns.isUnit(expressionType) && !KotlinBuiltIns.isNothing(expressionType)) return null
+
+                val context = statement.analyze()
+                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
+                }
                 return statement
             }
         }
diff --git a/idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt b/idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt
new file mode 100644
index 00000000000..e159d948b4f
--- /dev/null
+++ b/idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt
@@ -0,0 +1,9 @@
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(access: AccessMode) {
+    <caret>when (access) {
+        AccessMode.READ -> println("read")
+        AccessMode.WRITE -> println("write")
+        AccessMode.RW -> println("rw")
+    }
+}
+fun println(s: String) {}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt.after b/idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt.after
new file mode 100644
index 00000000000..41642dddd48
--- /dev/null
+++ b/idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt.after
@@ -0,0 +1,7 @@
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(access: AccessMode) = when (access) {
+    AccessMode.READ -> println("read")
+    AccessMode.WRITE -> println("write")
+    AccessMode.RW -> println("rw")
+}
+fun println(s: String) {}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToExpressionBody/whenUnitNonExhaustive.kt b/idea/testData/intentions/convertToExpressionBody/whenUnitNonExhaustive.kt
new file mode 100644
index 00000000000..25d33c8f8eb
--- /dev/null
+++ b/idea/testData/intentions/convertToExpressionBody/whenUnitNonExhaustive.kt
@@ -0,0 +1,9 @@
+// IS_APPLICABLE: false
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(access: AccessMode) {
+    <caret>when (access) {
+        AccessMode.READ -> println("read")
+        AccessMode.WRITE -> println("write")
+    }
+}
+fun println(s: String) {}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt b/idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt
new file mode 100644
index 00000000000..9029e255600
--- /dev/null
+++ b/idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt
@@ -0,0 +1,9 @@
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(access: AccessMode) {
+    <caret>when (access) {
+        AccessMode.READ -> println("read")
+        AccessMode.WRITE -> println("write")
+        else -> println("else")
+    }
+}
+fun println(s: String) {}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt.after b/idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt.after
new file mode 100644
index 00000000000..eb7cb5ef560
--- /dev/null
+++ b/idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt.after
@@ -0,0 +1,7 @@
+enum class AccessMode { READ, WRITE, RW }
+fun whenExpr(access: AccessMode) = when (access) {
+    AccessMode.READ -> println("read")
+    AccessMode.WRITE -> println("write")
+    else -> println("else")
+}
+fun println(s: String) {}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 5d9e71d8357..1151daa3c66 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -4591,6 +4591,24 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("whenUnitExhaustive.kt")
+        public void testWhenUnitExhaustive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/whenUnitExhaustive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("whenUnitNonExhaustive.kt")
+        public void testWhenUnitNonExhaustive() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/whenUnitNonExhaustive.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("whenUnitWithElse.kt")
+        public void testWhenUnitWithElse() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/whenUnitWithElse.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("while.kt")
         public void testWhile() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertToExpressionBody/while.kt");