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

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

@@ -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))
         }
     }
+
 }