Don't transform irBranch if tableswitch can't be generated

#KT-34466 Fixed

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -48740,6 +48740,18 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
             }
 
+            @Test
+            @TestMetadata("notAllConditionsAreEqEqExceptElse.kt")
+            public void testNotAllConditionsAreEqEqExceptElse() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/notAllConditionsAreEqEqExceptElse.kt");
+            }
+
+            @Test
+            @TestMetadata("notAllEntriesAreConstants.kt")
+            public void testNotAllEntriesAreConstants() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/notAllEntriesAreConstants.kt");
+            }
+
             @Test
             @TestMetadata("nullIsTheFirstEntry.kt")
             public void testNullIsTheFirstEntry() throws Exception {

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/EnumWhenLowering.kt

@@ -73,12 +73,39 @@ open class EnumWhenLowering(protected val context: CommonBackendContext) : IrEle
                 }
             }
         }
-
-        transformBranches(irWhen, subject, subjectOrdinalProvider)
-
+        if (possibleToGenerateJumpTable(irWhen, subject)) {
+            transformBranches(irWhen, subject, subjectOrdinalProvider)
+        }
         return expression
     }
 
+    private fun possibleToGenerateJumpTable(irWhen: IrWhen, subject: IrVariable): Boolean {
+        for (irBranch in irWhen.branches) {
+            val condition = irBranch.condition as? IrCall ?: continue
+            if (condition.symbol != context.irBuiltIns.eqeqSymbol)
+                return false
+
+            val lhs = condition.getValueArgument(0)!!
+            val rhs = condition.getValueArgument(1)!!
+            val other = getOther(lhs, rhs, subject)
+            if (other is IrCall) {
+                return false
+            }
+        }
+        return true
+    }
+
+    private fun getOther(lhs: IrExpression, rhs: IrExpression, subject: IrVariable): IrExpression? {
+        return when {
+            lhs is IrGetValue && lhs.symbol.owner == subject ->
+                rhs
+            rhs is IrGetValue && rhs.symbol.owner == subject ->
+                lhs
+            else ->
+                return null
+        }
+    }
+
     private fun transformBranches(
         irWhen: IrWhen,
         subject: IrVariable,
@@ -125,14 +152,7 @@ open class EnumWhenLowering(protected val context: CommonBackendContext) : IrEle
         }
         val lhs = expression.getValueArgument(0)!!
         val rhs = expression.getValueArgument(1)!!
-        val other = when {
-            lhs is IrGetValue && lhs.symbol.owner == subject ->
-                rhs
-            rhs is IrGetValue && rhs.symbol.owner == subject ->
-                lhs
-            else ->
-                return expression
-        }
+        val other = getOther(lhs, rhs, subject) ?: return expression
         val entryOrdinal = when {
             other is IrGetEnumValue && subject.type.classifierOrNull?.owner == other.symbol.owner.parent ->
                 mapConstEnumEntry(other.symbol.owner)

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

@@ -0,0 +1,25 @@
+// TARGET_BACKEND: JVM_IR
+// WITH_STDLIB
+
+enum class E {
+    A, B, C, D;
+}
+
+val abc = setOf(E.A, E.B, E.C)
+
+fun bar(): E = E.A
+
+fun foo(e: E): String {
+    val c = when (e) {
+        E.B -> "B"
+        in abc -> "OK"
+        else -> "else"
+    }
+    return c
+}
+
+fun box() = foo(bar())
+
+// CHECK_BYTECODE_TEXT
+// 0 WhenMappings
+// 0 TABLESWITCH

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

@@ -0,0 +1,22 @@
+// TARGET_BACKEND: JVM_IR
+
+enum class E {
+    A, B;
+}
+
+fun bar(): E = E.A
+
+fun foo(e: E): String {
+    val c = when (e) {
+        E.B -> "B"
+        bar() -> "OK"
+        else -> "else"
+    }
+    return c
+}
+
+fun box() = foo(bar())
+
+// CHECK_BYTECODE_TEXT
+// 0 WhenMappings
+// 0 TABLESWITCH

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

@@ -17,3 +17,7 @@ fun test(e: E?) = when (e) {
 fun box(): String {
     return test(null)
 }
+
+// CHECK_BYTECODE_TEXT
+// JVM_IR_TEMPLATES
+// 1 TABLESWITCH

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -48740,6 +48740,18 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
             }
 
+            @Test
+            @TestMetadata("notAllConditionsAreEqEqExceptElse.kt")
+            public void testNotAllConditionsAreEqEqExceptElse() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/notAllConditionsAreEqEqExceptElse.kt");
+            }
+
+            @Test
+            @TestMetadata("notAllEntriesAreConstants.kt")
+            public void testNotAllEntriesAreConstants() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/notAllEntriesAreConstants.kt");
+            }
+
             @Test
             @TestMetadata("nullIsTheFirstEntry.kt")
             public void testNullIsTheFirstEntry() throws Exception {