Turn off incorrect switch-optimization for when by enums

#KT-24708 Fixed

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

@@ -959,11 +959,18 @@ class CodegenAnnotatingVisitor extends KtVisitorVoid {
     }
 
     private boolean isWhenWithEnums(@NotNull KtWhenExpression expression) {
-        return WhenChecker.isWhenByEnum(expression, bindingContext) &&
+        ClassId enumClassId = WhenChecker.getClassIdForEnumSubject(expression, bindingContext);
-               switchCodegenProvider.checkAllItemsAreConstantsSatisfying(
+        if (enumClassId == null) return false;
-                       expression,
+
-                       constant -> constant instanceof EnumValue || constant instanceof NullValue
+        return switchCodegenProvider.checkAllItemsAreConstantsSatisfying(
-               );
+                expression,
+                constant -> isEnumEntryOrNull(enumClassId, constant)
+        );
+    }
+
+    private static boolean isEnumEntryOrNull(ClassId enumClassId, ConstantValue<?> constant) {
+        return (constant instanceof EnumValue && ((EnumValue) constant).getEnumClassId().equals(enumClassId)) ||
+               constant instanceof NullValue;
     }
 
     @NotNull

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

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils.isEnumClass
 import org.jetbrains.kotlin.resolve.DescriptorUtils.isEnumEntry
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
+import org.jetbrains.kotlin.resolve.descriptorUtil.classId
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeUtils
 import java.util.*
@@ -265,8 +266,12 @@ object WhenChecker {
     )
 
     @JvmStatic
-    fun isWhenByEnum(expression: KtWhenExpression, context: BindingContext) =
+    fun getClassIdForEnumSubject(expression: KtWhenExpression, context: BindingContext) =
-        getClassDescriptorOfTypeIfEnum(whenSubjectType(expression, context)) != null
+        getClassIdForTypeIfEnum(whenSubjectType(expression, context))
+
+    @JvmStatic
+    fun getClassIdForTypeIfEnum(type: KotlinType?) =
+        getClassDescriptorOfTypeIfEnum(type)?.classId
 
     @JvmStatic
     fun getClassDescriptorOfTypeIfEnum(type: KotlinType?): ClassDescriptor? {

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

@@ -0,0 +1,27 @@
+enum class A {
+    O, K
+}
+
+enum class B {
+    O, K
+}
+
+fun box(): String {
+    val a = A.O
+    val r1 = when (a) {
+        A.O -> "O"
+        A.K -> "K"
+        B.O -> "fail 1"
+        B.K -> "fail 2"
+    }
+
+    val b = B.K
+    val r2 = when (b) {
+        A.O -> "fail 3"
+        A.K -> "fail 4"
+        B.O -> "O"
+        B.K -> "K"
+    }
+
+    return r1 + r2
+}

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

@@ -0,0 +1,16 @@
+enum class A {
+    OK
+}
+
+enum class B {
+    FAIL
+}
+
+fun f() = A.OK
+
+fun box(): String {
+    return when (f()) {
+        B.FAIL -> "fail"
+        A.OK -> "OK"
+    }
+}

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

@@ -0,0 +1,30 @@
+enum class A {
+    O, K
+}
+
+enum class B {
+    O, K
+}
+
+fun box(): String {
+    val a = A.O
+    val r1 = when (a) {
+        A.O -> "O"
+        A.K -> "K"
+        B.O -> "fail 1"
+        B.K -> "fail 2"
+    }
+
+    val b = B.K
+    val r2 = when (b) {
+        A.O -> "fail 3"
+        A.K -> "fail 4"
+        B.O -> "O"
+        B.K -> "K"
+    }
+
+    return r1 + r2
+}
+
+// 0 TABLESWITCH
+// 0 LOOKUPSWITCH

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

@@ -0,0 +1,19 @@
+enum class A {
+    OK
+}
+
+enum class B {
+    FAIL
+}
+
+fun f() = A.OK
+
+fun box(): String {
+    return when (f()) {
+        B.FAIL -> "fail"
+        A.OK -> "OK"
+    }
+}
+
+// 0 TABLESWITCH
+// 0 LOOKUPSWITCH

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

@@ -22739,6 +22739,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/bigEnum.kt");
             }
 
+            @TestMetadata("differentEnumClasses.kt")
+            public void testDifferentEnumClasses() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses.kt");
+            }
+
+            @TestMetadata("differentEnumClasses2.kt")
+            public void testDifferentEnumClasses2() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses2.kt");
+            }
+
             @TestMetadata("duplicatingItems.kt")
             public void testDuplicatingItems() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/duplicatingItems.kt");

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

@@ -2962,6 +2962,16 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/bigEnum.kt");
         }
 
+        @TestMetadata("differentEnumClasses.kt")
+        public void testDifferentEnumClasses() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/differentEnumClasses.kt");
+        }
+
+        @TestMetadata("differentEnumClasses2.kt")
+        public void testDifferentEnumClasses2() throws Exception {
+            runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/differentEnumClasses2.kt");
+        }
+
         @TestMetadata("duplicatingItems.kt")
         public void testDuplicatingItems() throws Exception {
             runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/duplicatingItems.kt");

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

@@ -22739,6 +22739,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/when/enumOptimization/bigEnum.kt");
             }
 
+            @TestMetadata("differentEnumClasses.kt")
+            public void testDifferentEnumClasses() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses.kt");
+            }
+
+            @TestMetadata("differentEnumClasses2.kt")
+            public void testDifferentEnumClasses2() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses2.kt");
+            }
+
             @TestMetadata("duplicatingItems.kt")
             public void testDuplicatingItems() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/duplicatingItems.kt");

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

@@ -22739,6 +22739,16 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/when/enumOptimization/bigEnum.kt");
             }
 
+            @TestMetadata("differentEnumClasses.kt")
+            public void testDifferentEnumClasses() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses.kt");
+            }
+
+            @TestMetadata("differentEnumClasses2.kt")
+            public void testDifferentEnumClasses2() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses2.kt");
+            }
+
             @TestMetadata("duplicatingItems.kt")
             public void testDuplicatingItems() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/duplicatingItems.kt");

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

@@ -20554,6 +20554,16 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/bigEnum.kt");
             }
 
+            @TestMetadata("differentEnumClasses.kt")
+            public void testDifferentEnumClasses() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses.kt");
+            }
+
+            @TestMetadata("differentEnumClasses2.kt")
+            public void testDifferentEnumClasses2() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses2.kt");
+            }
+
             @TestMetadata("duplicatingItems.kt")
             public void testDuplicatingItems() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/duplicatingItems.kt");

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

@@ -21619,6 +21619,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/bigEnum.kt");
             }
 
+            @TestMetadata("differentEnumClasses.kt")
+            public void testDifferentEnumClasses() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses.kt");
+            }
+
+            @TestMetadata("differentEnumClasses2.kt")
+            public void testDifferentEnumClasses2() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/differentEnumClasses2.kt");
+            }
+
             @TestMetadata("duplicatingItems.kt")
             public void testDuplicatingItems() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/duplicatingItems.kt");

js/js.translator/src/org/jetbrains/kotlin/js/translate/expression/WhenTranslator.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.js.translate.expression
 
 import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.cfg.WhenChecker
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
@@ -208,11 +209,16 @@ private constructor(private val whenExpression: KtWhenExpression, context: Trans
             entries: List<KtWhenEntry>,
             expectedType: KotlinType
     ): Pair<List<EntryWithConstants>, Int> {
+        val classId = WhenChecker.getClassIdForTypeIfEnum(expectedType)
         return collectConstantEntries(
-                fromIndex, entries,
+            fromIndex, entries,
-                { (it.toConstantValue(expectedType) as? EnumValue)?.enumEntryName?.identifier },
+            {
-                { uniqueEnumNames.add(it) },
+                (it.toConstantValue(expectedType) as? EnumValue)
-                { JsStringLiteral(it) }
+                    ?.takeIf { enumEntry -> enumEntry.enumClassId == classId }
+                    ?.enumEntryName?.identifier
+            },
+            { uniqueEnumNames.add(it) },
+            { JsStringLiteral(it) }
         )
     }