Fix for KT-14597: "When" over smartcasted enum is broken and breaks all other "when"

#KT-14597 Fixed

compiler/backend/src/org/jetbrains/kotlin/codegen/binding/CodegenAnnotatingVisitor.java

@@ -648,8 +648,10 @@ class CodegenAnnotatingVisitor extends KtVisitorVoid {
         int fieldNumber = mappings.size();
 
         assert expression.getSubjectExpression() != null : "subject expression should be not null in a valid when by enums";
-        KotlinType type = bindingContext.getType(expression.getSubjectExpression());
+
+        KotlinType type = WhenChecker.whenSubjectType(expression, bindingContext);
         assert type != null : "should not be null in a valid when by enums";
+
         ClassDescriptor classDescriptor = (ClassDescriptor) type.getConstructor().getDeclarationDescriptor();
         assert classDescriptor != null : "because it's enum";
 

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

@@ -34,7 +34,7 @@ public class EnumSwitchCodegen extends SwitchCodegen {
             @NotNull ExpressionCodegen codegen,
             @NotNull WhenByEnumsMapping mapping
     ) {
-        super(expression, isStatement, isExhaustive, codegen);
+        super(expression, isStatement, isExhaustive, codegen, codegen.getState().getTypeMapper().mapType(mapping.getEnumClassDescriptor()));
         this.mapping = mapping;
     }
 

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

@@ -29,7 +29,7 @@ public class IntegralConstantsSwitchCodegen extends SwitchCodegen {
             boolean isExhaustive,
             @NotNull ExpressionCodegen codegen
     ) {
-        super(expression, isStatement, isExhaustive, codegen);
+        super(expression, isStatement, isExhaustive, codegen, null);
     }
 
     @Override

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

@@ -43,7 +43,7 @@ public class StringSwitchCodegen extends SwitchCodegen {
             boolean isExhaustive,
             @NotNull ExpressionCodegen codegen
     ) {
-        super(expression, isStatement, isExhaustive, codegen);
+        super(expression, isStatement, isExhaustive, codegen, null);
     }
 
     @Override

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

@@ -17,9 +17,9 @@
 package org.jetbrains.kotlin.codegen.when;
 
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.codegen.ExpressionCodegen;
 import org.jetbrains.kotlin.codegen.FrameMap;
-import org.jetbrains.kotlin.config.LanguageVersionSettings;
 import org.jetbrains.kotlin.psi.KtWhenEntry;
 import org.jetbrains.kotlin.psi.KtWhenExpression;
 import org.jetbrains.kotlin.resolve.BindingContext;
@@ -51,7 +51,8 @@ abstract public class SwitchCodegen {
 
     public SwitchCodegen(
             @NotNull KtWhenExpression expression, boolean isStatement,
-            boolean isExhaustive, @NotNull ExpressionCodegen codegen
+            boolean isExhaustive, @NotNull ExpressionCodegen codegen,
+            @Nullable Type subjectType
     ) {
         this.expression = expression;
         this.isStatement = isStatement;
@@ -59,7 +60,7 @@ abstract public class SwitchCodegen {
         this.codegen = codegen;
         this.bindingContext = codegen.getBindingContext();
 
-        subjectType = codegen.expressionType(expression.getSubjectExpression());
+        this.subjectType = subjectType != null ? subjectType : codegen.expressionType(expression.getSubjectExpression());
         resultType = isStatement ? Type.VOID_TYPE : codegen.expressionType(expression);
         v = codegen.v;
     }

compiler/frontend/src/org/jetbrains/kotlin/cfg/WhenChecker.kt

@@ -280,7 +280,8 @@ object WhenChecker {
         return classDescriptor
     }
 
-    private fun whenSubjectType(expression: KtWhenExpression, context: BindingContext) =
+    @JvmStatic
+    fun whenSubjectType(expression: KtWhenExpression, context: BindingContext) =
             expression.subjectExpression?.let { context.get(SMARTCAST, it)?.defaultType ?: context.getType(it) }
 
     @JvmStatic

compiler/testData/codegen/box/when/enumOptimization/kt14597.kt

@@ -0,0 +1,22 @@
+enum class En { A, B, С }
+
+fun box(): String {
+    var res = ""
+    // nullable variable
+    val en2: Any? = En.A
+    if (en2 is En) {
+        when (en2) {
+            En.A -> {res += "O"}
+            En.B -> {}
+            En.С -> {}
+        }
+
+        when (en2 as En) {
+            En.A -> {res += "K"}
+            En.B -> {}
+            En.С -> {}
+        }
+    }
+
+    return res
+}

compiler/testData/codegen/box/when/enumOptimization/kt14597_full.kt

@@ -0,0 +1,60 @@
+enum class En { A, B, С }
+
+fun main(args: Array<String>) {
+
+}
+
+fun box(): String {
+    var res1 = "fail"
+    var res2 = "fail2"
+
+    val en: En = En.A
+    when (en) {
+        En.A -> {res1 = ""}
+        En.B -> {}
+        En.С -> {}
+    }
+
+    when (en as En) {
+        En.A -> {res1 += "O"}
+        En.B -> {}
+        En.С -> {}
+    }
+
+
+    // nullable variable
+    val en2: Any? = En.A
+    if (en2 is En) {
+        when (en2) {
+            En.A -> {res1 += "K"}
+            En.B -> {}
+            En.С -> {}
+        }
+
+        when (en2 as En) {
+            En.A -> {res2 = ""}
+            En.B -> {}
+            En.С -> {}
+        }
+    }
+
+
+    // not nullable variable
+    val en1: Any = En.A
+    if (en1 is En) {
+        when (en1) {
+            En.A -> {res2 += "O"}
+            En.B -> {}
+            En.С -> {}
+        }
+        // Working without other examples
+        when (en1 as En) {
+            En.A -> {res2 += "K"}
+            En.B -> {}
+            En.С -> {}
+        }
+    }
+
+    if (res1 != res2) return "different results: $res1 != $res2"
+    return res1
+}

compiler/testData/codegen/box/when/enumOptimization/kt14802.kt

@@ -0,0 +1,25 @@
+class EncapsulatedEnum<T : Enum<T>>(val value: T)
+
+enum class MyEnum(val value: String) {
+    VALUE_A("OK"),
+    VALUE_B("fail"),
+}
+
+private fun crash(encapsulated: EncapsulatedEnum<*>) {
+    val myEnum = encapsulated.value
+    if (myEnum !is MyEnum) {
+        return
+    }
+
+    when (myEnum) {
+        MyEnum.VALUE_A -> res = myEnum.value
+        MyEnum.VALUE_B -> res = myEnum.value
+    }
+}
+
+var res = "fail"
+
+fun box(): String {
+    crash(EncapsulatedEnum(MyEnum.VALUE_A))
+    return res
+}

compiler/testData/codegen/box/when/enumOptimization/kt15806.kt

@@ -0,0 +1,25 @@
+
+private fun Any?.doTheThing(): String {
+    when (this) {
+        is String -> return this
+        is Level -> {
+            when (this) {
+                Level.O -> return Level.O.name
+                Level.K -> return Level.K.name
+            }
+        }
+
+        else -> return "fail"
+    }
+}
+
+
+enum class Level {
+    O,
+    K
+}
+
+
+fun box(): String {
+    return "O".doTheThing() + Level.K.doTheThing()
+}

compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14597_full.kt

@@ -0,0 +1,62 @@
+enum class En { A, B, С }
+
+fun main(args: Array<String>) {
+
+}
+
+fun box(): String {
+    var res1 = "fail"
+    var res2 = "fail2"
+
+    val en: En = En.A
+    when (en) {
+        En.A -> {res1 = ""}
+        En.B -> {}
+        En.С -> {}
+    }
+
+    when (en as En) {
+        En.A -> {res1 += "O"}
+        En.B -> {}
+        En.С -> {}
+    }
+
+
+    // nullable variable
+    val en2: Any? = En.A
+    if (en2 is En) {
+        when (en2) {
+            En.A -> {res1 += "K"}
+            En.B -> {}
+            En.С -> {}
+        }
+
+        when (en2 as En) {
+            En.A -> {res2 = ""}
+            En.B -> {}
+            En.С -> {}
+        }
+    }
+
+
+    // not nullable variable
+    val en1: Any = En.A
+    if (en1 is En) {
+        when (en1) {
+            En.A -> {res2 += "O"}
+            En.B -> {}
+            En.С -> {}
+        }
+        // Working without other examples
+        when (en1 as En) {
+            En.A -> {res2 += "K"}
+            En.B -> {}
+            En.С -> {}
+        }
+    }
+
+    if (res1 != res2) return "different results: $res1 != $res2"
+    return res1
+}
+
+// 6 TABLESWITCH

compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14802.kt

@@ -0,0 +1,27 @@
+class EncapsulatedEnum<T : Enum<T>>(val value: T)
+
+enum class MyEnum(val value: String) {
+    VALUE_A("OK"),
+    VALUE_B("fail"),
+}
+
+private fun crash(encapsulated: EncapsulatedEnum<*>) {
+    val myEnum = encapsulated.value
+    if (myEnum !is MyEnum) {
+        return
+    }
+
+    when (myEnum) {
+        MyEnum.VALUE_A -> res = myEnum.value
+        MyEnum.VALUE_B -> res = myEnum.value
+    }
+}
+
+var res = "fail"
+
+fun box(): String {
+    crash(EncapsulatedEnum(MyEnum.VALUE_A))
+    return res
+}
+
+// 1 LOOKUPSWITCH

compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt15806.kt

@@ -0,0 +1,27 @@
+
+private fun Any?.doTheThing(): String {
+    when (this) {
+        is String -> return this
+        is Level -> {
+            when (this) {
+                Level.O -> return Level.O.name
+                Level.K -> return Level.K.name
+            }
+        }
+
+        else -> return "fail"
+    }
+}
+
+
+enum class Level {
+    O,
+    K
+}
+
+
+fun box(): String {
+    return "O".doTheThing() + Level.K.doTheThing()
+}
+
+// 1 LOOKUPSWITCH

compiler/testData/codegen/light-analysis/when/enumOptimization/kt14597.txt

@@ -0,0 +1,14 @@
+@kotlin.Metadata
+public enum class En {
+    public final static field A: En
+    public final static field B: En
+    public final static field С: En
+    protected method <init>(p0: java.lang.String, p1: int): void
+    public static method valueOf(p0: java.lang.String): En
+    public static method values(): En[]
+}
+
+@kotlin.Metadata
+public final class Kt14597Kt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+}

compiler/testData/codegen/light-analysis/when/enumOptimization/kt14597_full.txt

@@ -0,0 +1,15 @@
+@kotlin.Metadata
+public enum class En {
+    public final static field A: En
+    public final static field B: En
+    public final static field С: En
+    protected method <init>(p0: java.lang.String, p1: int): void
+    public static method valueOf(p0: java.lang.String): En
+    public static method values(): En[]
+}
+
+@kotlin.Metadata
+public final class Kt14597_fullKt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method main(@org.jetbrains.annotations.NotNull p0: java.lang.String[]): void
+}

compiler/testData/codegen/light-analysis/when/enumOptimization/kt14802.txt

@@ -0,0 +1,26 @@
+@kotlin.Metadata
+public final class EncapsulatedEnum {
+    private final @org.jetbrains.annotations.NotNull field value: java.lang.Enum
+    public method <init>(@org.jetbrains.annotations.NotNull p0: java.lang.Enum): void
+    public final @org.jetbrains.annotations.NotNull method getValue(): java.lang.Enum
+}
+
+@kotlin.Metadata
+public final class Kt14802Kt {
+    private static @org.jetbrains.annotations.NotNull field res: java.lang.String
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    private final static method crash(p0: EncapsulatedEnum): void
+    public final static @org.jetbrains.annotations.NotNull method getRes(): java.lang.String
+    public final static method setRes(@org.jetbrains.annotations.NotNull p0: java.lang.String): void
+}
+
+@kotlin.Metadata
+public enum class MyEnum {
+    public final static field VALUE_A: MyEnum
+    public final static field VALUE_B: MyEnum
+    private final @org.jetbrains.annotations.NotNull field value: java.lang.String
+    protected method <init>(@java.lang.Synthetic p0: java.lang.String, @java.lang.Synthetic p1: int, @org.jetbrains.annotations.NotNull p2: java.lang.String): void
+    public final @org.jetbrains.annotations.NotNull method getValue(): java.lang.String
+    public static method valueOf(p0: java.lang.String): MyEnum
+    public static method values(): MyEnum[]
+}

compiler/testData/codegen/light-analysis/when/enumOptimization/kt15806.txt

@@ -0,0 +1,14 @@
+@kotlin.Metadata
+public final class Kt15806Kt {
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    private final static method doTheThing(@org.jetbrains.annotations.Nullable p0: java.lang.Object): java.lang.String
+}
+
+@kotlin.Metadata
+public enum class Level {
+    public final static field K: Level
+    public final static field O: Level
+    protected method <init>(p0: java.lang.String, p1: int): void
+    public static method valueOf(p0: java.lang.String): Level
+    public static method values(): Level[]
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -17689,6 +17689,30 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 doTest(fileName);
             }
 
+            @TestMetadata("kt14597.kt")
+            public void testKt14597() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14597_full.kt")
+            public void testKt14597_full() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597_full.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14802.kt")
+            public void testKt14802() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14802.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt15806.kt")
+            public void testKt15806() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("manyWhensWithinClass.kt")
             public void testManyWhensWithinClass() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");

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

@@ -17689,6 +17689,30 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("kt14597.kt")
+            public void testKt14597() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14597_full.kt")
+            public void testKt14597_full() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597_full.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14802.kt")
+            public void testKt14802() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14802.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt15806.kt")
+            public void testKt15806() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("manyWhensWithinClass.kt")
             public void testManyWhensWithinClass() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");

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

@@ -1712,6 +1712,24 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             doTest(fileName);
         }
 
+        @TestMetadata("kt14597_full.kt")
+        public void testKt14597_full() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14597_full.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("kt14802.kt")
+        public void testKt14802() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt14802.kt");
+            doTest(fileName);
+        }
+
+        @TestMetadata("kt15806.kt")
+        public void testKt15806() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/whenEnumOptimization/kt15806.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("manyWhensWithinClass.kt")
         public void testManyWhensWithinClass() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/whenEnumOptimization/manyWhensWithinClass.kt");

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

@@ -17689,6 +17689,30 @@ public class LightAnalysisModeCodegenTestGenerated extends AbstractLightAnalysis
                 doTest(fileName);
             }
 
+            @TestMetadata("kt14597.kt")
+            public void testKt14597() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14597_full.kt")
+            public void testKt14597_full() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597_full.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14802.kt")
+            public void testKt14802() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14802.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt15806.kt")
+            public void testKt15806() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("manyWhensWithinClass.kt")
             public void testManyWhensWithinClass() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -22313,6 +22313,30 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("kt14597.kt")
+            public void testKt14597() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14597_full.kt")
+            public void testKt14597_full() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14597_full.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt14802.kt")
+            public void testKt14802() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt14802.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("kt15806.kt")
+            public void testKt15806() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("manyWhensWithinClass.kt")
             public void testManyWhensWithinClass() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");