From 43b91589a1dc0b63c7cb3978314957799c7fa678 Mon Sep 17 00:00:00 2001
From: Pavel Kunyavskiy <pavel.kunyavskiy@jetbrains.com>
Date: Tue, 18 May 2021 14:16:07 +0300
Subject: [PATCH] [K/N] Create lowering for optimising when, with lowered enums
 support

---
 .../src/org/jetbrains/kotlin/backend/konan/Context.kt |  9 +++++++++
 .../kotlin/backend/konan/lower/EnumClassLowering.kt   | 11 ++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/Context.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/Context.kt
index ce778417859..bebd8b10cea 100644
--- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/Context.kt
+++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/Context.kt
@@ -104,6 +104,15 @@ internal class SpecialDeclarationsFactory(val context: Context) {
             }
         }
 
+    fun getLoweredEnumOrNull(enumClass: IrClass): LoweredEnumAccess? {
+        assert(enumClass.kind == ClassKind.ENUM_CLASS) { "Expected enum class but was: ${enumClass.descriptor}" }
+        return if (!context.llvmModuleSpecification.containsDeclaration(enumClass)) {
+            externalLoweredEnums[enumClass]
+        } else {
+            internalLoweredEnums[enumClass]
+        }
+    }
+
     fun getLoweredEnum(enumClass: IrClass): LoweredEnumAccess {
         assert(enumClass.kind == ClassKind.ENUM_CLASS) { "Expected enum class but was: ${enumClass.descriptor}" }
         return if (!context.llvmModuleSpecification.containsDeclaration(enumClass)) {
diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/EnumClassLowering.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/EnumClassLowering.kt
index de4509da01f..78c0623c439 100644
--- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/EnumClassLowering.kt
+++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/EnumClassLowering.kt
@@ -53,6 +53,15 @@ private class EnumSyntheticFunctionsBuilder(val context: Context) {
     private val genericValuesSymbol = context.ir.symbols.valuesForEnum
 }
 
+internal class NativeEnumWhenLowering constructor(context: Context) : EnumWhenLowering(context) {
+    override fun mapConstEnumEntry(entry: IrEnumEntry): Int {
+        val parent = entry.parentAsClass
+        val loweredEnum = (context as Context).specialDeclarationsFactory.getLoweredEnumOrNull(parent)
+                ?: return super.mapConstEnumEntry(entry)
+        return loweredEnum.entriesMap[entry.name]!!.ordinal
+    }
+}
+
 internal class EnumUsageLowering(val context: Context)
     : IrElementTransformer<IrBuilderWithScope?>, FileLoweringPass {
 
@@ -114,7 +123,7 @@ internal class EnumClassLowering(val context: Context) : FileLoweringPass {
     fun run(irFile: IrFile) {
         // EnumWhenLowering should be performed before EnumUsageLowering because
         // the latter performs lowering of IrGetEnumValue
-        EnumWhenLowering(context).lower(irFile)
+        NativeEnumWhenLowering(context).lower(irFile)
         lower(irFile)
         EnumUsageLowering(context).lower(irFile)
     }