Added switch (enum-only when) lowering

Added enum ordinal to serialization

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/Context.kt

@@ -51,10 +51,13 @@ import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
 import org.jetbrains.kotlin.konan.target.CompilerOutputKind
 import org.jetbrains.kotlin.load.java.BuiltinMethodsWithSpecialGenericSignature
+import org.jetbrains.kotlin.metadata.KonanLinkData
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitClassReceiver
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.getName
 import java.lang.System.out
 import kotlin.LazyThreadSafetyMode.PUBLICATION
 
@@ -63,6 +66,7 @@ internal class SpecialDeclarationsFactory(val context: Context) {
     private val outerThisFields = mutableMapOf<ClassDescriptor, IrField>()
     private val bridgesDescriptors = mutableMapOf<Pair<IrSimpleFunction, BridgeDirections>, IrSimpleFunction>()
     private val loweredEnums = mutableMapOf<ClassDescriptor, LoweredEnum>()
+    private val ordinals = mutableMapOf<IrClass, Map<ClassDescriptor, Int>>()
 
     object DECLARATION_ORIGIN_FIELD_FOR_OUTER_THIS :
             IrDeclarationOriginImpl("FIELD_FOR_OUTER_THIS")
@@ -123,6 +127,27 @@ internal class SpecialDeclarationsFactory(val context: Context) {
         }
     }
 
+    private fun assignOrdinalsToEnumEntries(irClass: IrClass): Map<ClassDescriptor, Int> {
+        val enumEntryOrdinals = mutableMapOf<ClassDescriptor, Int>()
+        irClass.declarations.filterIsInstance<IrEnumEntry>().forEachIndexed { index, entry ->
+            enumEntryOrdinals[entry.descriptor] = index
+        }
+        return enumEntryOrdinals
+    }
+
+    fun getEnumEntryOrdinal(entryDescriptor: ClassDescriptor): Int {
+        val enumClassDescriptor = entryDescriptor.containingDeclaration as ClassDescriptor
+        // If enum came from another module then we need to get serialized ordinal number.
+        // We serialize ordinal because current serialization cannot preserve enum entry order.
+        if (enumClassDescriptor is DeserializedClassDescriptor) {
+            return enumClassDescriptor.classProto.enumEntryList
+                    .first { entryDescriptor.name == enumClassDescriptor.c.nameResolver.getName(it.name) }
+                    .getExtension(KonanLinkData.enumEntryOrdinal)
+        }
+        val enumClass = context.ir.getEnum(enumClassDescriptor)
+        return ordinals.getOrPut(enumClass) { assignOrdinalsToEnumEntries(enumClass) }[entryDescriptor]!!
+    }
+
     private fun initializeBridgeDescriptor(bridgeDescriptor: SimpleFunctionDescriptorImpl,
                                            descriptor: FunctionDescriptor,
                                            bridgeDirections: Array<BridgeDirection>) {

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/KonanLower.kt

@@ -18,7 +18,6 @@ package org.jetbrains.kotlin.backend.konan
 
 import org.jetbrains.kotlin.backend.common.runOnFilePostfix
 import org.jetbrains.kotlin.backend.common.lower.*
-import org.jetbrains.kotlin.backend.common.validateIrModule
 import org.jetbrains.kotlin.backend.konan.irasdescriptors.referenceAllTypeExternalClassifiers
 import org.jetbrains.kotlin.backend.konan.lower.*
 import org.jetbrains.kotlin.backend.konan.lower.DefaultArgumentStubGenerator

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt

@@ -98,6 +98,7 @@ internal class KonanSymbols(context: Context, val symbolTable: SymbolTable): Sym
     val nothing = symbolTable.referenceClass(builtIns.nothing)
     val throwable = symbolTable.referenceClass(builtIns.throwable)
     val string = symbolTable.referenceClass(builtIns.string)
+    val enum = symbolTable.referenceClass(builtIns.enum)
 
     val arrayList = symbolTable.referenceClass(getArrayListClassDescriptor(context))
 

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/EnumClassLowering.kt

@@ -39,7 +39,6 @@ import org.jetbrains.kotlin.ir.declarations.impl.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
-import org.jetbrains.kotlin.ir.symbols.impl.IrAnonymousInitializerSymbolImpl
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
@@ -47,9 +46,8 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.scopes.MemberScope
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
-import org.jetbrains.kotlin.types.TypeProjectionImpl
+import org.jetbrains.kotlin.types.*
-import org.jetbrains.kotlin.types.TypeSubstitutor
+
-import org.jetbrains.kotlin.types.Variance
 
 internal class EnumSyntheticFunctionsBuilder(val context: Context) {
     fun buildValuesExpression(startOffset: Int, endOffset: Int,
@@ -159,8 +157,12 @@ internal class EnumUsageLowering(val context: Context)
 }
 
 internal class EnumClassLowering(val context: Context) : ClassLoweringPass {
+
     fun run(irFile: IrFile) {
         runOnFilePostfix(irFile)
+        // EnumWhenLowering should be performed before EnumUsageLowering because
+        // the latter performs lowering of IrGetEnumValue
+        EnumWhenLowering(context).lower(irFile)
         EnumUsageLowering(context).lower(irFile)
     }
 
@@ -177,7 +179,6 @@ internal class EnumClassLowering(val context: Context) : ClassLoweringPass {
 
     private inner class EnumClassTransformer(val irClass: IrClass) {
         private val loweredEnum = context.specialDeclarationsFactory.getLoweredEnum(irClass.descriptor)
-        private val enumEntryOrdinals = mutableMapOf<ClassDescriptor, Int>()
         private val loweredEnumConstructors = mutableMapOf<ClassConstructorDescriptor, IrConstructor>()
         private val descriptorToIrConstructorWithDefaultArguments = mutableMapOf<ClassConstructorDescriptor, IrConstructor>()
         private val defaultEnumEntryConstructors = mutableMapOf<ClassConstructorDescriptor, IrConstructor>()
@@ -186,7 +187,6 @@ internal class EnumClassLowering(val context: Context) : ClassLoweringPass {
 
         fun run() {
             insertInstanceInitializerCall()
-            assignOrdinalsToEnumEntries()
             lowerEnumConstructors(irClass)
             lowerEnumEntriesClasses()
             val defaultClass = createDefaultClassForEnumEntries()
@@ -221,16 +221,6 @@ internal class EnumClassLowering(val context: Context) : ClassLoweringPass {
             })
         }
 
-        private fun assignOrdinalsToEnumEntries() {
-            var ordinal = 0
-            irClass.declarations.forEach {
-                if (it is IrEnumEntry) {
-                    enumEntryOrdinals.put(it.descriptor, ordinal)
-                    ordinal++
-                }
-            }
-        }
-
         private fun lowerEnumEntriesClasses() {
             irClass.declarations.transformFlat { declaration ->
                 if (declaration is IrEnumEntry) {
@@ -563,7 +553,7 @@ internal class EnumClassLowering(val context: Context) : ClassLoweringPass {
         private abstract inner class InEnumEntry(private val enumEntry: ClassDescriptor) : EnumConstructorCallTransformer {
             override fun transform(enumConstructorCall: IrEnumConstructorCall): IrExpression {
                 val name = enumEntry.name.asString()
-                val ordinal = enumEntryOrdinals[enumEntry]!!
+                val ordinal = context.specialDeclarationsFactory.getEnumEntryOrdinal(enumEntry)
 
                 val descriptor = enumConstructorCall.descriptor
                 val startOffset = enumConstructorCall.startOffset

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

@@ -0,0 +1,134 @@
+package org.jetbrains.kotlin.backend.konan.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.peek
+import org.jetbrains.kotlin.backend.common.pop
+import org.jetbrains.kotlin.backend.common.push
+import org.jetbrains.kotlin.backend.konan.Context
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.declarations.impl.IrVariableImpl
+import org.jetbrains.kotlin.ir.descriptors.IrTemporaryVariableDescriptorImpl
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.util.getArguments
+import org.jetbrains.kotlin.ir.util.getPropertyGetter
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.types.isNullable
+
+/** Look for when-constructs where subject is enum entry.
+ * Replace branches that are comparisons with compile-time known enum entries
+ * with comparisons of ordinals.
+ */
+internal class EnumWhenLowering(private val context: Context) : IrElementTransformerVoid(), FileLoweringPass {
+
+    private val subjectWithOrdinalStack = mutableListOf<Pair<IrVariable, Lazy<IrVariable>>>()
+
+    override fun lower(irFile: IrFile) {
+        visitFile(irFile)
+    }
+
+    // Checks that irBlock satisfies all constrains of this lowering.
+    // 1. Block's origin is WHEN
+    // 2. Subject of `when` is variable of enum type
+    // NB: See BranchingExpressionGenerator in Kotlin sources to get insight about
+    // `when` block translation to IR.
+    private fun shouldLower(irBlock: IrBlock): Boolean {
+        if (irBlock.origin != IrStatementOrigin.WHEN) {
+            return false
+        }
+        // when-block with subject should have two children: temporary variable and when itself.
+        if (irBlock.statements.size != 2) {
+            return false
+        }
+        val subject = irBlock.statements[0] as IrVariable
+        // Subject should not be nullable because we will access the `ordinal` property.
+        if (subject.type.isNullable()) {
+            return false
+        }
+        // Check that subject is enum entry.
+        val enumClass = subject.type.constructor.declarationDescriptor as? ClassDescriptor
+                ?: return false
+        return enumClass.kind == ClassKind.ENUM_CLASS
+    }
+
+    override fun visitBlock(expression: IrBlock): IrExpression {
+        if (!shouldLower(expression)) {
+            return super.visitBlock(expression)
+        }
+        // Will be initialized only when we found a branch that compares
+        // subject with compile-time known enum entry.
+        val subject = expression.statements[0] as IrVariable
+        val subjectOrdinalProvider = lazy {
+            createEnumOrdinalVariable(subject)
+        }
+        subjectWithOrdinalStack.push(Pair(subject, subjectOrdinalProvider))
+        // Process nested `when` and comparisons.
+        expression.transformChildrenVoid(this)
+        // If variable was initialized then it was actually used and we need to insert it
+        // into the block's IR.
+        if (subjectOrdinalProvider.isInitialized()) {
+            expression.statements.add(1, subjectOrdinalProvider.value)
+        }
+        subjectWithOrdinalStack.pop()
+        return expression
+    }
+
+    private fun createEnumOrdinalVariable(enumVariable: IrVariable): IrVariable {
+        val ordinalPropertyGetter = context.ir.symbols.enum.getPropertyGetter("ordinal")!!
+        val getOrdinal = IrCallImpl(enumVariable.startOffset, enumVariable.endOffset, ordinalPropertyGetter).apply {
+            dispatchReceiver = IrGetValueImpl(enumVariable.startOffset, enumVariable.endOffset, enumVariable.symbol)
+        }
+        // Create temporary variable for subject's ordinal.
+        val ordinalDescriptor = IrTemporaryVariableDescriptorImpl(enumVariable.descriptor.containingDeclaration,
+                Name.identifier(enumVariable.name.asString() + "_ordinal"), context.builtIns.intType)
+        return IrVariableImpl(enumVariable.startOffset, enumVariable.endOffset,
+                IrDeclarationOrigin.IR_TEMPORARY_VARIABLE, ordinalDescriptor, getOrdinal)
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        expression.transformChildrenVoid(this)
+        // Try to do actual lowering.
+        return tryLower(expression)
+    }
+
+    private val areEqualByValue = context.ir.symbols.areEqualByValue.first {
+        it.owner.valueParameters[0].type == context.builtIns.intType
+    }
+
+    // We are looking for branch that is a comparison of the subject and another enum entry.
+    private fun tryLower(call: IrCall): IrExpression {
+        if (call.origin != IrStatementOrigin.EQEQ) {
+            return call
+        }
+        val callArgs = call.getArguments()
+        if (callArgs.size != 2) {
+            return call
+        }
+        val lhs = callArgs[0].second
+        val rhs = callArgs[1].second
+        // Both entries should belong to the same class.
+        if (lhs.type != rhs.type) {
+            return call
+        }
+        // If there is nothing on stack then nothing we can do.
+        val (topmostSubject, topmostOrdinalProvider) = subjectWithOrdinalStack.peek()
+                ?: return call
+        if (lhs is IrValueAccessExpression && lhs.symbol.owner == topmostSubject && rhs is IrGetEnumValue) {
+            val entryOrdinal = context.specialDeclarationsFactory.getEnumEntryOrdinal(rhs.descriptor)
+            val subjectOrdinal = topmostOrdinalProvider.value
+            return IrCallImpl(call.startOffset, call.endOffset, areEqualByValue).apply {
+                putValueArgument(0, IrGetValueImpl(lhs.startOffset, lhs.endOffset, subjectOrdinal.symbol))
+                putValueArgument(1, IrConstImpl.int(rhs.startOffset, rhs.endOffset, context.builtIns.intType, entryOrdinal))
+            }
+        }
+        return call
+    }
+}

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/serialization/KonanLinkData.proto

@@ -44,6 +44,7 @@ extend Property {
 
 extend EnumEntry {
     repeated Annotation enum_entry_annotation = 170;
+    optional int32 enum_entry_ordinal = 171;
 }
 
 extend ValueParameter {

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/serialization/KonanSerializerExtension.kt

@@ -48,7 +48,9 @@ internal class KonanSerializerExtension(val context: Context) :
     }
 
     override fun serializeEnumEntry(descriptor: ClassDescriptor, proto: ProtoBuf.EnumEntry.Builder) {
-
+        // Serialization doesn't preserve enum entry order, so we need to serialize ordinal.
+        val ordinal = context.specialDeclarationsFactory.getEnumEntryOrdinal(descriptor)
+        proto.setExtension(KonanLinkData.enumEntryOrdinal, ordinal)
         super.serializeEnumEntry(descriptor, proto)
     }
 

backend.native/tests/build.gradle

@@ -763,6 +763,11 @@ task enum_loop(type: RunKonanTest) {
     source = "codegen/enum/loop.kt"
 }
 
+task switchLowering(type: RunKonanTest) {
+    goldValue = "EnumA.A\nok\nok\nok\nok\nok\n"
+    source = "codegen/enum/switchLowering.kt"
+}
+
 task mangling(type: LinkKonanTest) {
     goldValue = 
         "Int direct [1, 2, 3, 4]\n" +
@@ -2420,6 +2425,12 @@ task serialized_no_typemap(type: RunStandaloneKonanTest) {
     goldValue = "OK\n"
 }
 
+task serialized_enum_ordinal(type: LinkKonanTest) {
+    source = "serialization/enum_ordinal/main.kt"
+    lib = "serialization/enum_ordinal/library.kt"
+    goldValue = "0\n1\n2\nb\n"
+}
+
 task testing_annotations(type: RunStandaloneKonanTest) {
     source = "testing/annotations.kt"
     flags = ['-tr']

backend.native/tests/codegen/enum/switchLowering.kt

@@ -0,0 +1,86 @@
+package codegen.enum.switchLowering
+
+import kotlin.test.*
+
+enum class EnumA {
+    A, B, C
+}
+
+enum class EnumB {
+    A, B
+}
+
+enum class E {
+    ONE, TWO, THREE
+}
+
+fun produceEntry() = EnumA.A
+
+// Check that we fail on comparison of different enum types.
+fun differentEnums() {
+    println(when (produceEntry()) {
+        EnumB.A -> "EnumB.A"
+        EnumA.A -> "EnumA.A"
+        EnumA.B -> "EnumA.B"
+        else    -> "nah"
+    })
+}
+
+// Nullable subject shouldn't be lowered.
+fun nullable() {
+    val x: EnumA? = null
+    when(x) {
+        EnumA.A -> println("fail")
+        else    -> println("ok")
+    }
+}
+
+// Operator overloading won't trick us!
+fun operatorOverloading() {
+    operator fun E.contains(other: E): Boolean = false
+
+    val y = E.ONE
+    when(y) {
+        in E.ONE    -> println("Should not reach here")
+        else        -> println("ok")
+    }
+}
+
+fun smoke1() {
+    when (produceEntry()) {
+        EnumA.B -> println("error")
+        EnumA.A -> println("ok")
+        EnumA.C -> println("error")
+    }
+}
+
+fun smoke2() {
+    when (produceEntry()) {
+        EnumA.B -> println("error")
+        else    -> println("ok")
+    }
+}
+
+fun eA() = EnumA.A
+
+fun eB() = EnumA.B
+
+
+fun nestedWhen() {
+    println(when (eA()) {
+        EnumA.A, EnumA.C -> when (eB()) {
+            EnumA.B -> "ok"
+            else -> "nope"
+        }
+        else -> "nope"
+    })
+}
+
+@Test fun runTest() {
+    differentEnums()
+    nullable()
+    operatorOverloading()
+    smoke1()
+    smoke2()
+    nestedWhen()
+}

backend.native/tests/serialization/enum_ordinal/library.kt

@@ -0,0 +1,9 @@
+enum class Color {
+    RED, GREEN, BLUE, CYAN, MAGENTA, YELLOW
+}
+
+fun determineColor(code: Int): Color = when (code) {
+    0 -> Color.BLUE
+    1 -> Color.MAGENTA
+    else -> Color.CYAN
+}

backend.native/tests/serialization/enum_ordinal/main.kt

@@ -0,0 +1,14 @@
+fun main(args: Array<String>) {
+    println(Color.RED.ordinal)
+    println(Color.GREEN.ordinal)
+    println(Color.BLUE.ordinal)
+    val color = when (determineColor(args.size)) {
+        Color.RED -> println("r")
+        Color.GREEN -> println("g")
+        Color.BLUE -> println("b")
+        Color.CYAN -> println("c")
+        Color.MAGENTA -> println("m")
+        Color.YELLOW -> println("y")
+
+    }
+}

runtime/src/main/kotlin/konan/internal/RuntimeUtils.kt

@@ -94,8 +94,7 @@ fun ReportUnhandledException(e: Throwable) {
 @ExportForCppRuntime
 internal fun TheEmptyString() = ""
 
-fun <T: Enum<T>> valueOfForEnum(name: String, values: Array<T>) : T
+fun <T: Enum<T>> valueOfForEnum(name: String, values: Array<T>) : T {
-{
     var left = 0
     var right = values.size - 1
     while (left <= right) {
@@ -110,8 +109,7 @@ fun <T: Enum<T>> valueOfForEnum(name: String, values: Array<T>) : T
     throw Exception("Invalid enum name: $name")
 }
 
-fun <T: Enum<T>> valuesForEnum(values: Array<T>): Array<T>
+fun <T: Enum<T>> valuesForEnum(values: Array<T>): Array<T> {
-{
     val result = @Suppress("TYPE_PARAMETER_AS_REIFIED") Array<T?>(values.size)
     for (value in values)
         result[value.ordinal] = value