IR KT-50258 rewrite enum-based 'when' more carefully

This commit is contained in:
Dmitry Petrov
2021-12-21 06:06:45 +00:00
committed by Space
parent e17f121b23
commit 8905586cbb
16 changed files with 308 additions and 28 deletions
@@ -47838,12 +47838,30 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@Test
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@Test
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -5886,6 +5886,12 @@ public class FirBytecodeTextTestGenerated extends AbstractFirBytecodeTextTest {
runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -5,7 +5,9 @@
package org.jetbrains.kotlin.backend.common.lower
import org.jetbrains.kotlin.backend.common.*
import org.jetbrains.kotlin.backend.common.CommonBackendContext
import org.jetbrains.kotlin.backend.common.FileLoweringPass
import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
import org.jetbrains.kotlin.descriptors.ClassKind
import org.jetbrains.kotlin.ir.builders.*
import org.jetbrains.kotlin.ir.declarations.IrEnumEntry
@@ -19,14 +21,12 @@ import org.jetbrains.kotlin.ir.types.classifierOrNull
import org.jetbrains.kotlin.ir.types.getClass
import org.jetbrains.kotlin.ir.types.isNullable
import org.jetbrains.kotlin.ir.util.*
import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
/** Look for when-constructs where subject is enum entry.
* Replace branches that are comparisons with compile-time known enum entries
* with comparisons of ordinals.
*/
open class EnumWhenLowering(protected val context: CommonBackendContext) : IrElementTransformerVoidWithContext(), FileLoweringPass {
private val subjectWithOrdinalStack = mutableListOf<Pair<IrVariable, Lazy<IrVariable>>>()
protected open fun mapConstEnumEntry(entry: IrEnumEntry): Int =
entry.parentAsClass.declarations.filterIsInstance<IrEnumEntry>().indexOf(entry).also {
@@ -41,22 +41,24 @@ open class EnumWhenLowering(protected val context: CommonBackendContext) : IrEle
}
override fun visitBlock(expression: IrBlock): IrExpression {
expression.transformChildrenVoid()
// NB: See BranchingExpressionGenerator to get insight about `when` block translation to IR.
if (expression.origin != IrStatementOrigin.WHEN) {
return super.visitBlock(expression)
return expression
}
// when-block with subject should have two children: temporary variable and when itself.
if (expression.statements.size != 2) {
return super.visitBlock(expression)
}
val subject = expression.statements[0]
if (subject !is IrVariable) {
return super.visitBlock(expression)
return expression
}
val subject = expression.statements[0] as? IrVariable
?: return expression
val subjectClass = subject.type.getClass()
if (subjectClass == null || subjectClass.kind != ClassKind.ENUM_CLASS || subjectClass.isEffectivelyExternal()) {
return super.visitBlock(expression)
return expression
}
val irWhen = expression.statements[1] as? IrWhen
?: return expression
// Will be initialized only when we found a branch that compares
// subject with compile-time known enum entry.
@@ -71,39 +73,75 @@ open class EnumWhenLowering(protected val context: CommonBackendContext) : IrEle
}
}
}
subjectWithOrdinalStack.push(Pair(subject, subjectOrdinalProvider))
try {
// Process nested `when` and comparisons.
expression.statements[1].transformChildrenVoid(this)
} finally {
subjectWithOrdinalStack.pop()
}
transformBranches(irWhen, subject, subjectOrdinalProvider)
return expression
}
override fun visitCall(expression: IrCall): IrExpression {
private fun transformBranches(
irWhen: IrWhen,
subject: IrVariable,
subjectOrdinalProvider: Lazy<IrVariable>
): IrExpression {
for (irBranch in irWhen.branches) {
irBranch.condition = transformBranchSubexpression(irBranch.condition, subject, subjectOrdinalProvider)
irBranch.result = transformBranchSubexpression(irBranch.result, subject, subjectOrdinalProvider)
}
return irWhen
}
private fun transformBranchSubexpression(
irExpression: IrExpression,
subject: IrVariable,
subjectOrdinalProvider: Lazy<IrVariable>
): IrExpression =
when (irExpression) {
is IrCall -> {
// Single entry clause:
// when (subject) {
// ENUM_ENTRY -> ...
// }
transformEnumEquals(irExpression, subject, subjectOrdinalProvider)
}
is IrWhen -> {
// Multiple entry clause:
// when (subject) {
// ENUM_ENTRY_1, ENUM_ENTRY_2, ..., ENUM_ENTRY_N -> ...
// }
transformBranches(irExpression, subject, subjectOrdinalProvider)
}
else -> irExpression
}
private fun transformEnumEquals(
expression: IrCall,
subject: IrVariable,
subjectOrdinalProvider: Lazy<IrVariable>
): IrExpression {
// We are looking for branch that is a comparison of the subject and another enum entry.
if (expression.symbol != context.irBuiltIns.eqeqSymbol) {
return super.visitCall(expression)
return expression
}
val lhs = expression.getValueArgument(0)!!
val rhs = expression.getValueArgument(1)!!
val (topmostSubject, topmostOrdinalProvider) = subjectWithOrdinalStack.peek()
?: return super.visitCall(expression)
val other = when {
lhs is IrGetValue && lhs.symbol.owner == topmostSubject -> rhs
rhs is IrGetValue && rhs.symbol.owner == topmostSubject -> lhs
else -> return super.visitCall(expression)
lhs is IrGetValue && lhs.symbol.owner == subject ->
rhs
rhs is IrGetValue && rhs.symbol.owner == subject ->
lhs
else ->
return expression
}
val entryOrdinal = when {
other is IrGetEnumValue && topmostSubject.type.classifierOrNull?.owner == other.symbol.owner.parent ->
other is IrGetEnumValue && subject.type.classifierOrNull?.owner == other.symbol.owner.parent ->
mapConstEnumEntry(other.symbol.owner)
other.isNullConst() ->
-1
else -> return super.visitCall(expression)
else ->
return expression
}
val subjectOrdinal = topmostOrdinalProvider.value
val subjectOrdinal = subjectOrdinalProvider.value
return IrCallImpl(
expression.startOffset, expression.endOffset,
expression.type, expression.symbol,
@@ -114,4 +152,5 @@ open class EnumWhenLowering(protected val context: CommonBackendContext) : IrEle
putValueArgument(1, IrConstImpl.int(rhs.startOffset, rhs.endOffset, context.irBuiltIns.intType, entryOrdinal))
}
}
}
@@ -0,0 +1,23 @@
enum class MyEnum {
First,
Second
}
fun getValue() = MyEnum.First
var result = "Failed"
fun getLambda(): (Int) -> Unit =
when (val value = getValue()) {
MyEnum.Second -> { _ -> }
MyEnum.First -> { _ ->
if (value == MyEnum.First) {
result = "OK"
}
}
}
fun box(): String {
getLambda().invoke(2)
return result
}
@@ -0,0 +1,19 @@
enum class ABCD {
A, B, C, D
}
fun test(x: ABCD, y: ABCD, ok: String): String =
when (x) {
when (y) {
ABCD.A -> ABCD.B
ABCD.B -> ABCD.C
ABCD.C -> ABCD.D
ABCD.D -> ABCD.A
} ->
ok
ABCD.A, ABCD.B, ABCD.C, ABCD.D ->
x.toString()
}
fun box(): String =
test(ABCD.B, ABCD.A, "O") + test(ABCD.A, ABCD.D, "K")
@@ -0,0 +1,21 @@
enum class ABCD {
A, B, C, D
}
fun test(x: ABCD, y: ABCD, ok: String): String =
when (x) {
ABCD.A, ABCD.B ->
when (y) {
ABCD.A, ABCD.B -> ok
ABCD.C, ABCD.D -> y.toString()
}
ABCD.C, ABCD.D ->
x.toString()
}
fun box(): String =
test(ABCD.B, ABCD.A, "O") + test(ABCD.A, ABCD.B, "K")
// CHECK_BYTECODE_TEXT
// JVM_IR_TEMPLATES
// 2 (TABLE|LOOKUP)SWITCH
@@ -0,0 +1,22 @@
// IGNORE_BACKEND_FIR: JVM_IR
// FIR_STATUS: FIR+JVM_IR generates TABLESWITCH for nested 'when' only, doesn't look critical.
// JVM_IR_TEMPLATES
// 2 (TABLE|LOOKUP)SWITCH
enum class ABCD {
A, B, C, D
}
fun test(x: ABCD, y: ABCD, ok: String): String =
when (x) {
when (y) {
ABCD.A -> ABCD.B
ABCD.B -> ABCD.C
ABCD.C -> ABCD.D
ABCD.D -> ABCD.A
} ->
ok
ABCD.A, ABCD.B, ABCD.C, ABCD.D ->
x.toString()
}
@@ -45084,12 +45084,30 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@Test
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@Test
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -5742,6 +5742,12 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -47838,12 +47838,30 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@Test
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@Test
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -5886,6 +5886,12 @@ public class IrBytecodeTextTestGenerated extends AbstractIrBytecodeTextTest {
runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/bytecodeText/whenEnumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -36598,11 +36598,26 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
@@ -32756,12 +32756,30 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@Test
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@Test
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -32858,12 +32858,30 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@Test
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@Test
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
@@ -27522,11 +27522,26 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
@@ -35169,12 +35169,30 @@ public class ExternalTestGenerated extends AbstractExternalNativeBlackBoxTest {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt15806.kt");
}
@Test
@TestMetadata("kt50258.kt")
public void testKt50258() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/kt50258.kt");
}
@Test
@TestMetadata("manyWhensWithinClass.kt")
public void testManyWhensWithinClass() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/manyWhensWithinClass.kt");
}
@Test
@TestMetadata("nestedWhenInCondition.kt")
public void testNestedWhenInCondition() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInCondition.kt");
}
@Test
@TestMetadata("nestedWhenInResult.kt")
public void testNestedWhenInResult() throws Exception {
runTest("compiler/testData/codegen/box/when/enumOptimization/nestedWhenInResult.kt");
}
@Test
@TestMetadata("nonConstantEnum.kt")
public void testNonConstantEnum() throws Exception {