diff --git a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java b/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java index 8f560a2c571..56892b4c0de 100644 --- a/compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java +++ b/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 { diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/EnumWhenLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/EnumWhenLowering.kt index 01cbd423bc6..f432b8fea59 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/EnumWhenLowering.kt +++ b/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) diff --git a/compiler/testData/codegen/box/when/enumOptimization/notAllConditionsAreEqEqExceptElse.kt b/compiler/testData/codegen/box/when/enumOptimization/notAllConditionsAreEqEqExceptElse.kt new file mode 100644 index 00000000000..7976ed6cb1a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/compiler/testData/codegen/box/when/enumOptimization/notAllEntriesAreConstants.kt b/compiler/testData/codegen/box/when/enumOptimization/notAllEntriesAreConstants.kt new file mode 100644 index 00000000000..ab28ac66b88 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/compiler/testData/codegen/box/when/enumOptimization/nullableEnum.kt b/compiler/testData/codegen/box/when/enumOptimization/nullableEnum.kt index bfce330f0bc..7707e796750 100644 --- a/compiler/testData/codegen/box/when/enumOptimization/nullableEnum.kt +++ b/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 diff --git a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java b/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java index 183fc9998a3..a5a0eaeaef9 100644 --- a/compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java +++ b/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 {