JVM IR: fix "null" as the first entry in MappedEnumWhenLowering

If "null" was the first entry in an optimizable "when" over enum,
mapRuntimeEnumEntry was called before mapConstEnumEntry, and the
$WhenMappings field was not created. Now both mapConstEnumEntry and
mapRuntimeEnumEntry create this field on the first access

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/MappedEnumWhenLowering.kt

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetEnumValueImpl
 import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
@@ -69,7 +68,7 @@ private class MappedEnumWhenLowering(context: CommonBackendContext) : EnumWhenLo
     // of the classes in which they are used. This field tracks which container is the innermost one.
     private var state: EnumMappingState? = null
 
-    private class EnumMappingState {
+    private inner class EnumMappingState {
         val mappings = mutableMapOf<IrClass /* enum */, Pair<MutableMap<IrEnumEntry, Int>, IrField>>()
         val mappingsClass by lazy {
             buildClass {
@@ -79,26 +78,28 @@ private class MappedEnumWhenLowering(context: CommonBackendContext) : EnumWhenLo
                 createImplicitParameterDeclarationWithWrappedDescriptor()
             }
         }
+
+        fun getMappingForClass(enumClass: IrClass): Pair<MutableMap<IrEnumEntry, Int>, IrField> =
+            mappings.getOrPut(enumClass) {
+                mutableMapOf<IrEnumEntry, Int>() to mappingsClass.addField {
+                    name = Name.identifier("\$EnumSwitchMapping\$${mappings.size}")
+                    type = intArray.owner.defaultType
+                    origin = JvmLoweredDeclarationOrigin.ENUM_MAPPINGS_FOR_WHEN
+                    isFinal = true
+                    isStatic = true
+                }
+            }
     }
 
     override fun mapConstEnumEntry(entry: IrEnumEntry): Int {
-        val (mapping, _) = state!!.mappings.getOrPut(entry.parentAsClass) {
-            mutableMapOf<IrEnumEntry, Int>() to state!!.mappingsClass.addField {
-                name = Name.identifier("\$EnumSwitchMapping\$${state!!.mappings.size}")
-                type = intArray.owner.defaultType
-                origin = JvmLoweredDeclarationOrigin.ENUM_MAPPINGS_FOR_WHEN
-                isFinal = true
-                isStatic = true
-            }
-        }
+        val (mapping) = state!!.getMappingForClass(entry.parentAsClass)
         // Index 0 (default value for integers) is reserved for unknown ordinals.
         return mapping.getOrPut(entry) { mapping.size + 1 }
     }
 
     override fun mapRuntimeEnumEntry(builder: IrBuilderWithScope, subject: IrExpression): IrExpression =
         builder.irCall(intArrayGet).apply {
-            val (_, field) = state!!.mappings[subject.type.getClass()!!]
-                ?: throw AssertionError("no values mapped for enum class ${subject.type}")
+            val (_, field) = state!!.getMappingForClass(subject.type.getClass()!!)
             dispatchReceiver = builder.irGetField(null, field)
             putValueArgument(0, super.mapRuntimeEnumEntry(builder, subject))
         }

compiler/testData/codegen/box/when/enumOptimization/nullIsTheFirstEntry.kt

@@ -0,0 +1,14 @@
+enum class E {
+    A, B;
+}
+
+fun foo(e: E?): String {
+    val c = when (e) {
+        null -> "Fail: null"
+        E.B -> "OK"
+        E.A -> "Fail: A"
+    }
+    return c
+}
+
+fun box(): String = foo(E.B)

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -26192,6 +26192,11 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
             }
 
+            @TestMetadata("nullIsTheFirstEntry.kt")
+            public void testNullIsTheFirstEntry() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/nullIsTheFirstEntry.kt");
+            }
+
             @TestMetadata("nullability.kt")
             public void testNullability() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nullability.kt");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -26192,6 +26192,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
             }
 
+            @TestMetadata("nullIsTheFirstEntry.kt")
+            public void testNullIsTheFirstEntry() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/nullIsTheFirstEntry.kt");
+            }
+
             @TestMetadata("nullability.kt")
             public void testNullability() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nullability.kt");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -25087,6 +25087,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
             }
 
+            @TestMetadata("nullIsTheFirstEntry.kt")
+            public void testNullIsTheFirstEntry() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/nullIsTheFirstEntry.kt");
+            }
+
             @TestMetadata("nullability.kt")
             public void testNullability() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nullability.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -20172,6 +20172,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
             }
 
+            @TestMetadata("nullIsTheFirstEntry.kt")
+            public void testNullIsTheFirstEntry() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/nullIsTheFirstEntry.kt");
+            }
+
             @TestMetadata("nullability.kt")
             public void testNullability() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nullability.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -21327,6 +21327,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nonConstantEnum.kt");
             }
 
+            @TestMetadata("nullIsTheFirstEntry.kt")
+            public void testNullIsTheFirstEntry() throws Exception {
+                runTest("compiler/testData/codegen/box/when/enumOptimization/nullIsTheFirstEntry.kt");
+            }
+
             @TestMetadata("nullability.kt")
             public void testNullability() throws Exception {
                 runTest("compiler/testData/codegen/box/when/enumOptimization/nullability.kt");