Implicit exhaustive whens now have exception in else branch #KT-8700 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -3801,8 +3801,9 @@ The "returned" value of try expression with no finally is either the last expres
         return StackValue.operation(resultType, new Function1<InstructionAdapter, Unit>() {
             @Override
             public Unit invoke(InstructionAdapter v) {
-                SwitchCodegen switchCodegen =
-                        SwitchCodegenUtil.buildAppropriateSwitchCodegenIfPossible(expression, isStatement, ExpressionCodegen.this);
+                SwitchCodegen switchCodegen = SwitchCodegenUtil.buildAppropriateSwitchCodegenIfPossible(
+                        expression, isStatement, isExhaustive(expression, isStatement), ExpressionCodegen.this
+                );
                 if (switchCodegen != null) {
                     switchCodegen.generate();
                     return Unit.INSTANCE;
@@ -3848,9 +3849,7 @@ The "returned" value of try expression with no finally is either the last expres
                 }
                 if (!hasElse && nextCondition != null) {
                     v.mark(nextCondition);
-                    if (!isStatement) {
-                        putUnitInstanceOntoStackForNonExhaustiveWhen(expression);
-                    }
+                    putUnitInstanceOntoStackForNonExhaustiveWhen(expression, isStatement);
                 }
 
                 markLineNumber(expression, isStatement);
@@ -3863,14 +3862,24 @@ The "returned" value of try expression with no finally is either the last expres
         });
     }
 
+    private boolean isExhaustive(@NotNull KtWhenExpression whenExpression, boolean isStatement) {
+        if (isStatement) {
+            return Boolean.TRUE.equals(bindingContext.get(BindingContext.IMPLICIT_EXHAUSTIVE_WHEN, whenExpression));
+        }
+        else {
+            return Boolean.TRUE.equals(bindingContext.get(BindingContext.EXHAUSTIVE_WHEN, whenExpression));
+        }
+    }
+
     public void putUnitInstanceOntoStackForNonExhaustiveWhen(
-            @NotNull KtWhenExpression expression
+            @NotNull KtWhenExpression whenExpression,
+            boolean isStatement
     ) {
-        if (Boolean.TRUE.equals(bindingContext.get(BindingContext.EXHAUSTIVE_WHEN, expression))) {
+        if (isExhaustive(whenExpression, isStatement)) {
             // when() is supposed to be exhaustive
             genThrow(v, "kotlin/NoWhenBranchMatchedException", null);
         }
-        else {
+        else if (!isStatement) {
             // non-exhaustive when() with no else -> Unit must be expected
             StackValue.putUnitInstance(v);
         }

compiler/backend/src/org/jetbrains/kotlin/codegen/when/EnumSwitchCodegen.java

@@ -30,10 +30,11 @@ public class EnumSwitchCodegen extends SwitchCodegen {
     public EnumSwitchCodegen(
             @NotNull KtWhenExpression expression,
             boolean isStatement,
+            boolean isExhaustive,
             @NotNull ExpressionCodegen codegen,
             @NotNull WhenByEnumsMapping mapping
     ) {
-        super(expression, isStatement, codegen);
+        super(expression, isStatement, isExhaustive, codegen);
         this.mapping = mapping;
     }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/when/IntegralConstantsSwitchCodegen.java

@@ -26,9 +26,10 @@ public class IntegralConstantsSwitchCodegen extends SwitchCodegen {
     public IntegralConstantsSwitchCodegen(
             @NotNull KtWhenExpression expression,
             boolean isStatement,
+            boolean isExhaustive,
             @NotNull ExpressionCodegen codegen
     ) {
-        super(expression, isStatement, codegen);
+        super(expression, isStatement, isExhaustive, codegen);
     }
 
     @Override

compiler/backend/src/org/jetbrains/kotlin/codegen/when/StringSwitchCodegen.java

@@ -40,9 +40,10 @@ public class StringSwitchCodegen extends SwitchCodegen {
     public StringSwitchCodegen(
             @NotNull KtWhenExpression expression,
             boolean isStatement,
+            boolean isExhaustive,
             @NotNull ExpressionCodegen codegen
     ) {
-        super(expression, isStatement, codegen);
+        super(expression, isStatement, isExhaustive, codegen);
     }
 
     @Override

compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegen.java

@@ -35,6 +35,7 @@ import java.util.*;
 abstract public class SwitchCodegen {
     protected final KtWhenExpression expression;
     protected final boolean isStatement;
+    protected final boolean isExhaustive;
     protected final ExpressionCodegen codegen;
     protected final BindingContext bindingContext;
     protected final Type subjectType;
@@ -49,10 +50,11 @@ abstract public class SwitchCodegen {
 
     public SwitchCodegen(
             @NotNull KtWhenExpression expression, boolean isStatement,
-            @NotNull ExpressionCodegen codegen
+            boolean isExhaustive, @NotNull ExpressionCodegen codegen
     ) {
         this.expression = expression;
         this.isStatement = isStatement;
+        this.isExhaustive = isExhaustive;
         this.codegen = codegen;
         this.bindingContext = codegen.getBindingContext();
 
@@ -70,7 +72,7 @@ abstract public class SwitchCodegen {
         boolean hasElse = expression.getElseExpression() != null;
 
         // if there is no else-entry and it's statement then default --- endLabel
-        defaultLabel = (hasElse || !isStatement) ? elseLabel : endLabel;
+        defaultLabel = (hasElse || !isStatement || isExhaustive) ? elseLabel : endLabel;
 
         generateSubject();
 
@@ -79,9 +81,9 @@ abstract public class SwitchCodegen {
         generateEntries();
 
         // there is no else-entry but this is not statement, so we should return Unit
-        if (!hasElse && !isStatement) {
+        if (!hasElse && (!isStatement || isExhaustive)) {
             v.visitLabel(elseLabel);
-            codegen.putUnitInstanceOntoStackForNonExhaustiveWhen(expression);
+            codegen.putUnitInstanceOntoStackForNonExhaustiveWhen(expression, isStatement);
         }
 
         codegen.markLineNumber(expression, isStatement);

compiler/backend/src/org/jetbrains/kotlin/codegen/when/SwitchCodegenUtil.java

@@ -102,6 +102,7 @@ public class SwitchCodegenUtil {
     public static SwitchCodegen buildAppropriateSwitchCodegenIfPossible(
             @NotNull KtWhenExpression expression,
             boolean isStatement,
+            boolean isExhaustive,
             @NotNull ExpressionCodegen codegen
     ) {
         BindingContext bindingContext = codegen.getBindingContext();
@@ -114,15 +115,15 @@ public class SwitchCodegenUtil {
         WhenByEnumsMapping mapping = codegen.getBindingContext().get(CodegenBinding.MAPPING_FOR_WHEN_BY_ENUM, expression);
 
         if (mapping != null) {
-            return new EnumSwitchCodegen(expression, isStatement, codegen, mapping);
+            return new EnumSwitchCodegen(expression, isStatement, isExhaustive, codegen, mapping);
         }
 
         if (isIntegralConstantsSwitch(expression, subjectType, bindingContext)) {
-            return new IntegralConstantsSwitchCodegen(expression, isStatement, codegen);
+            return new IntegralConstantsSwitchCodegen(expression, isStatement, isExhaustive, codegen);
         }
 
         if (isStringConstantsSwitch(expression, subjectType, bindingContext)) {
-            return new StringSwitchCodegen(expression, isStatement, codegen);
+            return new StringSwitchCodegen(expression, isStatement, isExhaustive, codegen);
         }
 
         return null;

compiler/testData/codegen/box/when/exhaustiveWhenInitialization.kt

@@ -0,0 +1,10 @@
+enum class A { V }
+
+fun box(): String {
+  val a: A = A.V
+  val b: Boolean
+  when (a) {
+    A.V -> b = true
+  }
+  return if (b) "OK" else "FAIL"
+}

compiler/testData/codegen/box/when/exhaustiveWhenReturn.kt

@@ -0,0 +1,8 @@
+enum class A { V }
+
+fun box(): String {
+    val a: A = A.V
+    when (a) {
+        A.V -> return "OK"
+    }
+}

compiler/testData/codegen/box/when/sealedWhenInitialization.kt

@@ -0,0 +1,15 @@
+sealed class A {
+    object B : A()
+
+    class C : A()
+}
+
+fun box(): String {
+    val a: A = A.C()
+    val b: Boolean
+    when (a) {
+        A.B -> b = true
+        is A.C -> b = false
+    }
+    return if (!b) "OK" else "FAIL"
+}

compiler/testData/codegen/boxWithStdlib/when/exhaustiveBreakContinue.kt

@@ -0,0 +1,14 @@
+enum class Color { RED, GREEN, BLUE }
+
+fun foo(arr: Array<Color>): Color {
+    loop@ for (color in arr) {
+        when (color) {
+            Color.RED -> return color
+            Color.GREEN -> break@loop
+            Color.BLUE -> if (arr.size == 1) return color else continue@loop
+        }
+    }
+    return Color.GREEN
+}
+
+fun box() = if (foo(arrayOf(Color.BLUE, Color.GREEN)) == Color.GREEN) "OK" else "FAIL"

compiler/testData/codegen/bytecodeText/when/exhaustiveWhenInitialization.kt

@@ -0,0 +1,14 @@
+enum class A { V }
+
+fun box(): String {
+  val a: A = A.V
+  val b: Boolean
+  when (a) {
+    A.V -> b = true
+  }
+  return if (b) "OK" else "FAIL"
+}
+
+// 0 TABLESWITCH
+// 1 LOOKUPSWITCH
+// 1 ATHROW

compiler/testData/codegen/bytecodeText/when/exhaustiveWhenReturn.kt

@@ -0,0 +1,12 @@
+enum class A { V }
+
+fun box(): String {
+    val a: A = A.V
+    when (a) {
+        A.V -> return "OK"
+    }
+}
+
+// 0 TABLESWITCH
+// 1 LOOKUPSWITCH
+// 1 ATHROW

compiler/testData/codegen/bytecodeText/when/sealedWhenInitialization.kt

@@ -0,0 +1,19 @@
+sealed class A {
+    object B : A()
+
+    class C : A()
+}
+
+fun box(): String {
+    val a: A = A.C()
+    val b: Boolean
+    when (a) {
+        A.B -> b = true
+        is A.C -> b = false
+    }
+    return if (!b) "OK" else "FAIL"
+}
+
+// 0 TABLESWITCH
+// 0 LOOKUPSWITCH
+// 1 ATHROW

compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java

@@ -1039,11 +1039,29 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/bytecodeText/when"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("exhaustiveWhenInitialization.kt")
+        public void testExhaustiveWhenInitialization() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/when/exhaustiveWhenInitialization.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("exhaustiveWhenReturn.kt")
+        public void testExhaustiveWhenReturn() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/when/exhaustiveWhenReturn.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("integralWhenWithNoInlinedConstants.kt")
         public void testIntegralWhenWithNoInlinedConstants() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/when/integralWhenWithNoInlinedConstants.kt");
             doTest(fileName);
         }
+
+        @TestMetadata("sealedWhenInitialization.kt")
+        public void testSealedWhenInitialization() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/when/sealedWhenInitialization.kt");
+            doTest(fileName);
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/bytecodeText/whenEnumOptimization")

compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxCodegenTestGenerated.java

@@ -8254,6 +8254,18 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("exhaustiveWhenInitialization.kt")
+        public void testExhaustiveWhenInitialization() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/exhaustiveWhenInitialization.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("exhaustiveWhenReturn.kt")
+        public void testExhaustiveWhenReturn() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/exhaustiveWhenReturn.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("is.kt")
         public void testIs() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/is.kt");
@@ -8332,6 +8344,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("sealedWhenInitialization.kt")
+        public void testSealedWhenInitialization() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/sealedWhenInitialization.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("whenArgumentIsEvaluatedOnlyOnce.kt")
         public void testWhenArgumentIsEvaluatedOnlyOnce() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/whenArgumentIsEvaluatedOnlyOnce.kt");

compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxWithStdlibCodegenTestGenerated.java

@@ -4891,6 +4891,12 @@ public class BlackBoxWithStdlibCodegenTestGenerated extends AbstractBlackBoxCode
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxWithStdlib/when"), Pattern.compile("^(.+)\\.kt$"), true);
         }
 
+        @TestMetadata("exhaustiveBreakContinue.kt")
+        public void testExhaustiveBreakContinue() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxWithStdlib/when/exhaustiveBreakContinue.kt");
+            doTestWithStdlib(fileName);
+        }
+
         @TestMetadata("integralWhenWithNoInlinedConstants.kt")
         public void testIntegralWhenWithNoInlinedConstants() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/boxWithStdlib/when/integralWhenWithNoInlinedConstants.kt");