JVM: support enumEntries intrinsic for Java & old Kotlin enums

#KT-59710 Fixed

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/EnumEntriesIntrinsicMappingsCacheImpl.kt

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.jvm.codegen
+
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.lower.irBlockBody
+import org.jetbrains.kotlin.backend.jvm.EnumEntriesIntrinsicMappingsCache
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.backend.jvm.ir.findEnumValuesFunction
+import org.jetbrains.kotlin.ir.builders.declarations.addField
+import org.jetbrains.kotlin.ir.builders.declarations.addFunction
+import org.jetbrains.kotlin.ir.builders.declarations.buildClass
+import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.builders.irSetField
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrField
+import org.jetbrains.kotlin.ir.types.typeWith
+import org.jetbrains.kotlin.ir.util.createImplicitParameterDeclarationWithWrappedDescriptor
+import org.jetbrains.kotlin.ir.util.defaultType
+import org.jetbrains.kotlin.load.java.JavaDescriptorVisibilities
+import org.jetbrains.kotlin.name.Name
+
+// This class generates synthetic `$EntriesIntrinsicMappings` classes which cache the result of calling the `enumEntries` intrinsic
+// on Java enums and on old (pre-1.9) Kotlin enums.
+//
+// `$EntriesIntrinsicMappings` classes are exactly the same as `$EntriesMappings` classes generated in `EnumExternalEntriesLowering`
+// Unfortunately, we cannot reuse the logic from that lowering as is because, as opposed to the normal function call `E.entries`,
+// the call to the intrinsic `enumEntries<E>()` is fully reified only during the codegen, where all IR of the module has already been
+// lowered, notably:
+// 1) All the `$EntriesMappings` have already been created for the `E.entries` calls, so we cannot easily reuse those classes, which is the
+//    the reason why we're generating new classes with a slightly different name, even though it might be suboptimal in cases when you have
+//    both `E.entries` and `enumEntries<E>` calls to the same enum in the same container.
+// 2) Static initializers have been lowered, so we're generating the `<clinit>` method manually, as opposed to adding static init sections
+//    as `EnumExternalEntriesLowering` does.
+class EnumEntriesIntrinsicMappingsCacheImpl(private val context: JvmBackendContext) : EnumEntriesIntrinsicMappingsCache() {
+    private val storage = mutableMapOf<IrClass, MappingsClass>()
+
+    private inner class MappingsClass(val containingClass: IrClass) {
+        val irClass = context.irFactory.buildClass {
+            name = Name.identifier("EntriesIntrinsicMappings")
+            origin = JvmLoweredDeclarationOrigin.ENUM_MAPPINGS_FOR_ENTRIES
+        }.apply {
+            createImplicitParameterDeclarationWithWrappedDescriptor()
+            parent = containingClass
+        }
+
+        val enums = hashMapOf<IrClass, IrField>()
+    }
+
+    @Synchronized
+    override fun getEnumEntriesIntrinsicMappings(containingClass: IrClass, enumClass: IrClass): IrField {
+        val mappingsClass = storage.getOrPut(containingClass) { MappingsClass(containingClass) }
+        val field = mappingsClass.enums.getOrPut(enumClass) {
+            mappingsClass.irClass.addField {
+                name = Name.identifier("entries\$${mappingsClass.enums.size}")
+                type = context.ir.symbols.enumEntries.typeWith(enumClass.defaultType)
+                origin = JvmLoweredDeclarationOrigin.ENUM_MAPPINGS_FOR_ENTRIES
+                isFinal = true
+                isStatic = true
+            }
+        }
+        return field
+    }
+
+    override fun generateMappingsClasses() {
+        val backendContext = context
+        for (klass in storage.values) {
+            // Use the same origin and visibility for `<clinit>` as in StaticInitializersLowering.
+            val clinit = klass.irClass.addFunction(
+                "<clinit>", context.irBuiltIns.unitType, visibility = JavaDescriptorVisibilities.PACKAGE_VISIBILITY,
+                isStatic = true, origin = JvmLoweredDeclarationOrigin.CLASS_STATIC_INITIALIZER
+            )
+            clinit.body = context.createIrBuilder(clinit.symbol).irBlockBody(clinit) {
+                // Sort fields by name. Note that if there are a lot of calls to entries of different enums in the same container, it would
+                // result in the ordering "entries$0, entries$1, entries$10, entries$11, entries$12, ...", but it's not a big deal.
+                for ((enum, field) in klass.enums.entries.sortedBy { it.value.name }) {
+                    // For each field, we're generating:
+                    //     entries$N = kotlin.enums.EnumEntriesKt.enumEntries(E.values())
+                    val enumValues = enum.findEnumValuesFunction(backendContext)
+                    +irSetField(
+                        null, field,
+                        irCall(backendContext.ir.symbols.createEnumEntries).apply {
+                            putValueArgument(0, irCall(enumValues))
+                        }
+                    )
+                }
+            }
+
+            ClassCodegen.getOrCreate(klass.irClass, backendContext).generate()
+        }
+    }
+}

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineIntrinsicsSupport.kt

@@ -33,6 +33,7 @@ import org.jetbrains.org.objectweb.asm.Type.INT_TYPE
 import org.jetbrains.org.objectweb.asm.Type.VOID_TYPE
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter
 import org.jetbrains.org.objectweb.asm.tree.AbstractInsnNode
+import org.jetbrains.org.objectweb.asm.tree.FieldInsnNode
 import org.jetbrains.org.objectweb.asm.tree.InsnList
 
 class IrInlineIntrinsicsSupport(
@@ -119,6 +120,9 @@ class IrInlineIntrinsicsSupport(
 
     override fun toKotlinType(type: IrType): KotlinType = type.toIrBasedKotlinType()
 
+    override fun generateExternalEntriesForEnumTypeIfNeeded(type: IrType): FieldInsnNode? =
+        generateExternalEntriesForEnumTypeIfNeeded(type, classCodegen)
+
     override fun reportSuspendTypeUnsupported() {
         classCodegen.context.ktDiagnosticReporter.at(reportErrorsOn, containingFile).report(JvmBackendErrors.TYPEOF_SUSPEND_TYPE)
     }

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/codegen/irCodegenUtils.kt

@@ -38,6 +38,7 @@ import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmClassSignature
 import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.org.objectweb.asm.Opcodes
 import org.jetbrains.org.objectweb.asm.Type
+import org.jetbrains.org.objectweb.asm.tree.FieldInsnNode
 import kotlin.collections.set
 
 class IrFrameMap : FrameMapBase<IrSymbol>() {
@@ -349,3 +350,15 @@ val IrClass.reifiedTypeParameters: ReifiedTypeParametersUsages
 
         return tempReifiedTypeParametersUsages
     }
+
+internal fun generateExternalEntriesForEnumTypeIfNeeded(type: IrType, containingCodegen: ClassCodegen): FieldInsnNode? {
+    val irClass = type.getClass()
+    if (irClass == null || !irClass.isEnumClassWhichRequiresExternalEntries()) return null
+
+    val mappingsCache = containingCodegen.context.enumEntriesIntrinsicMappingsCache
+    val field = mappingsCache.getEnumEntriesIntrinsicMappings(containingCodegen.irClass, irClass)
+    return FieldInsnNode(
+        Opcodes.GETSTATIC, containingCodegen.typeMapper.mapClass(field.parentAsClass).internalName,
+        field.name.asString(), AsmTypes.ENUM_ENTRIES.descriptor
+    )
+}

compiler/ir/backend.jvm/codegen/src/org/jetbrains/kotlin/backend/jvm/intrinsics/EnumIntrinsics.kt

@@ -80,8 +80,12 @@ object EnumEntries : IntrinsicMethod() {
             mv.checkcast(AsmTypes.ENUM_ENTRIES)
             MaterialValue(codegen, AsmTypes.ENUM_ENTRIES, expression.type)
         } else {
-            // TODO: support enumEntries for Java enums.
-            mv.invokestatic(typeMapper.mapType(type).internalName, "getEntries", Type.getMethodDescriptor(AsmTypes.ENUM_ENTRIES), false)
+            val getField = generateExternalEntriesForEnumTypeIfNeeded(type, codegen.classCodegen)
+            if (getField != null) {
+                mv.visitFieldInsn(getField.opcode, getField.owner, getField.name, getField.desc)
+            } else {
+                mv.invokestatic(typeMapper.mapType(type).internalName, "getEntries", Type.getMethodDescriptor(AsmTypes.ENUM_ENTRIES), false)
+            }
             MaterialValue(codegen, AsmTypes.ENUM_ENTRIES, expression.type)
         }
     }

compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.backend.common.phaser.CompilerPhase
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
 import org.jetbrains.kotlin.backend.common.serialization.DescriptorByIdSignatureFinderImpl
+import org.jetbrains.kotlin.backend.jvm.codegen.EnumEntriesIntrinsicMappingsCacheImpl
 import org.jetbrains.kotlin.backend.jvm.codegen.JvmIrIntrinsicExtension
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
 import org.jetbrains.kotlin.backend.jvm.ir.getIoFile
@@ -343,6 +344,9 @@ open class JvmIrCodegenFactory(
         context.getIntrinsic = { symbol: IrFunctionSymbol ->
             intrinsics.getIntrinsic(symbol) ?: generationExtensions.firstNotNullOfOrNull { it.getIntrinsic(symbol) }
         }
+
+        context.enumEntriesIntrinsicMappingsCache = EnumEntriesIntrinsicMappingsCacheImpl(context)
+
         /* JvmBackendContext creates new unbound symbols, have to resolve them. */
         ExternalDependenciesGenerator(symbolTable, irProviders).generateUnboundSymbolsAsDependencies()
 

compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmPhases.kt

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.backend.jvm
 
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.phaser.SameTypeNamedCompilerPhase
-import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
-import org.jetbrains.kotlin.backend.common.phaser.performByIrFile
-import org.jetbrains.kotlin.backend.common.phaser.then
+import org.jetbrains.kotlin.backend.common.phaser.*
 import org.jetbrains.kotlin.backend.jvm.codegen.ClassCodegen
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrFile
@@ -47,6 +44,22 @@ private class FileCodegen(private val context: JvmBackendContext, private val ge
     }
 }
 
+private val generateAdditionalClassesPhase = SameTypeNamedCompilerPhase(
+    "GenerateAdditionalClasses",
+    "Generate additional classes that were requested during codegen",
+    lower = object : SameTypeCompilerPhase<JvmBackendContext, IrModuleFragment> {
+        override fun invoke(
+            phaseConfig: PhaseConfigurationService,
+            phaserState: PhaserState<IrModuleFragment>,
+            context: JvmBackendContext,
+            input: IrModuleFragment,
+        ): IrModuleFragment {
+            context.enumEntriesIntrinsicMappingsCache.generateMappingsClasses()
+            return input
+        }
+    }
+)
+
 // Generate multifile facades first, to compute and store JVM signatures of const properties which are later used
 // when serializing metadata in the multifile parts.
 // TODO: consider dividing codegen itself into separate phases (bytecode generation, metadata serialization) to avoid this
@@ -55,7 +68,8 @@ internal val jvmCodegenPhases = SameTypeNamedCompilerPhase(
     description = "Code generation",
     nlevels = 1,
     lower = codegenPhase(generateMultifileFacade = true) then
-            codegenPhase(generateMultifileFacade = false)
+            codegenPhase(generateMultifileFacade = false) then
+            generateAdditionalClassesPhase
 )
 
 // This property is needed to avoid dependencies from "leaf" modules (cli, tests-common-new) on backend.jvm:lower.

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

@@ -11,7 +11,8 @@ import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
-import org.jetbrains.kotlin.backend.jvm.ir.isInCurrentModule
+import org.jetbrains.kotlin.backend.jvm.ir.findEnumValuesFunction
+import org.jetbrains.kotlin.backend.jvm.ir.isEnumClassWhichRequiresExternalEntries
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.builders.declarations.addField
@@ -26,7 +27,8 @@ import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
 import org.jetbrains.kotlin.ir.types.typeWith
-import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.ir.util.createImplicitParameterDeclarationWithWrappedDescriptor
+import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
@@ -60,6 +62,8 @@ internal val enumExternalEntriesPhase = makeIrFilePhase(
  * // F.kt
  * FKt$EntriesMappings.entries$1
  * ```
+ *
+ * There's similar code which handles the `enumEntries<Enum>()` intrinsic code generation in `EnumEntriesIntrinsicMappingsCacheImpl`.
  */
 class EnumExternalEntriesLowering(private val context: JvmBackendContext) : FileLoweringPass, IrElementTransformerVoidWithContext() {
 
@@ -99,33 +103,12 @@ class EnumExternalEntriesLowering(private val context: JvmBackendContext) : File
     override fun visitCall(expression: IrCall): IrExpression {
         val owner = expression.symbol.owner as? IrSimpleFunction
         val parentClass = owner?.parent as? IrClass ?: return super.visitCall(expression)
-        /*
-         * Candidates for lowering:
-         * * Java enums
-         * * Kotlin enums that have no 'getEntries' function (thus compiled with pre-1.8 LV/AV)
-         */
-        val shouldBeLowered = parentClass.isEnumClass &&
-                owner.name == SpecialNames.ENUM_GET_ENTRIES &&
-                (parentClass.isFromJava() || !parentClass.hasEnumEntriesFunction())
+        val shouldBeLowered = owner.name == SpecialNames.ENUM_GET_ENTRIES && parentClass.isEnumClassWhichRequiresExternalEntries()
         if (!shouldBeLowered) return super.visitCall(expression)
         val field = state!!.getEntriesFieldForEnum(parentClass)
         return IrGetFieldImpl(expression.startOffset, expression.endOffset, field.symbol, field.type)
     }
 
-    private fun IrClass.hasEnumEntriesFunction(): Boolean {
-        // Enums from other modules are always loaded with a property `entries` which has a getter `<get-entries>`.
-        // Enums from the current module will have a property `entries` if they are unlowered yet (i.e. enum is declared in another file
-        // which will be lowered after the file with the call site), or a function `<get-entries>` if they are already lowered.
-        return functions.any { it.isGetEntries() }
-                || (properties.any { it.getter?.isGetEntries() == true } && isInCurrentModule())
-    }
-
-    private fun IrSimpleFunction.isGetEntries(): Boolean =
-        name.toString() == "<get-entries>"
-                && dispatchReceiverParameter == null
-                && extensionReceiverParameter == null
-                && valueParameters.isEmpty()
-
     override fun visitClassNew(declaration: IrClass): IrStatement {
         val oldState = state
         val mappingState = EntriesMappingState()

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.backend.common.lower.irCatch
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
+import org.jetbrains.kotlin.backend.jvm.ir.findEnumValuesFunction
 import org.jetbrains.kotlin.backend.jvm.ir.isInPublicInlineScope
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
@@ -23,7 +24,8 @@ import org.jetbrains.kotlin.ir.declarations.IrEnumEntry
 import org.jetbrains.kotlin.ir.declarations.IrField
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetEnumValueImpl
-import org.jetbrains.kotlin.ir.types.*
+import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -167,12 +169,3 @@ private class MappedEnumWhenLowering(override val context: JvmBackendContext) :
         return declaration
     }
 }
-
-internal fun IrClass.findEnumValuesFunction(context: JvmBackendContext) = functions.single {
-    it.name.toString() == "values"
-            && it.dispatchReceiverParameter == null
-            && it.extensionReceiverParameter == null
-            && it.valueParameters.isEmpty()
-            && it.returnType.isBoxedArray
-            && it.returnType.getArrayElementType(context.irBuiltIns).classOrNull == this.symbol
-}

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/EnumEntriesIntrinsicMappingsCache.kt

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.backend.jvm
+
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrField
+
+abstract class EnumEntriesIntrinsicMappingsCache {
+    abstract fun getEnumEntriesIntrinsicMappings(containingClass: IrClass, enumClass: IrClass): IrField
+
+    abstract fun generateMappingsClasses()
+}

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -114,6 +114,8 @@ class JvmBackendContext(
 
     lateinit var getIntrinsic: (IrFunctionSymbol) -> IntrinsicMarker?
 
+    lateinit var enumEntriesIntrinsicMappingsCache: EnumEntriesIntrinsicMappingsCache
+
     // Store evaluated SMAP for anonymous classes. Used only with IR inliner.
     val typeToCachedSMAP = mutableMapOf<Type, SMAP>()
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/ir/JvmIrUtils.kt

@@ -518,3 +518,30 @@ fun IrFunction.extensionReceiverName(state: GenerationState): String {
 
 fun IrFunction.isBridge(): Boolean =
     origin == IrDeclarationOrigin.BRIDGE || origin == IrDeclarationOrigin.BRIDGE_SPECIAL
+
+// Enum requires external implementation of entries if it's either a Java enum, or a Kotlin enum compiled with pre-1.8 LV/AV.
+fun IrClass.isEnumClassWhichRequiresExternalEntries(): Boolean =
+    isEnumClass && (isFromJava() || !hasEnumEntriesFunction())
+
+private fun IrClass.hasEnumEntriesFunction(): Boolean {
+    // Enums from other modules are always loaded with a property `entries` which has a getter `<get-entries>`.
+    // Enums from the current module will have a property `entries` if they are unlowered yet (i.e. enum is declared in another file
+    // which will be lowered after the file with the call site), or a function `<get-entries>` if they are already lowered.
+    return functions.any { it.isGetEntries() }
+            || (properties.any { it.getter?.isGetEntries() == true } && isInCurrentModule())
+}
+
+private fun IrSimpleFunction.isGetEntries(): Boolean =
+    name.toString() == "<get-entries>"
+            && dispatchReceiverParameter == null
+            && extensionReceiverParameter == null
+            && valueParameters.isEmpty()
+
+fun IrClass.findEnumValuesFunction(context: JvmBackendContext): IrSimpleFunction = functions.single {
+    it.name.toString() == "values"
+            && it.dispatchReceiverParameter == null
+            && it.extensionReceiverParameter == null
+            && it.valueParameters.isEmpty()
+            && it.returnType.isBoxedArray
+            && it.returnType.getArrayElementType(context.irBuiltIns).classOrNull == this.symbol
+}