[FIR2IR] Part 4. Cleanup Fir2IrClassifierStorage API (enum entries)

^KT-65937 Fixed

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrClassifierStorage.kt

@@ -41,7 +41,7 @@ class Fir2IrClassifierStorage(
 
     private val typeParameterCacheForSetter: MutableMap<FirTypeParameter, IrTypeParameter> = mutableMapOf()
 
-    private val enumEntryCache: MutableMap<FirEnumEntry, IrEnumEntry> = commonMemberStorage.enumEntryCache
+    private val enumEntryCache: MutableMap<FirEnumEntry, IrEnumEntrySymbol> = commonMemberStorage.enumEntryCache
 
     private val codeFragmentCache: MutableMap<FirCodeFragment, IrClass> = mutableMapOf()
 
@@ -64,7 +64,7 @@ class Fir2IrClassifierStorage(
     fun forEachCachedDeclarationSymbol(block: (IrSymbol) -> Unit) {
         classCache.values.forEach { block(it) }
         typeAliasCache.values.forEach { block(it.symbol) }
-        enumEntryCache.values.forEach { block(it.symbol) }
+        enumEntryCache.values.forEach { block(it) }
         fieldsForContextReceivers.values.forEach { fields ->
             fields.forEach { block(it.symbol) }
         }
@@ -332,17 +332,30 @@ class Fir2IrClassifierStorage(
         localStorage[(enumEntry.initializer as FirAnonymousObjectExpression).anonymousObject] = correspondingClass
     }
 
-    internal fun getCachedIrEnumEntry(enumEntry: FirEnumEntry): IrEnumEntry? {
+    fun getIrEnumEntrySymbol(enumEntry: FirEnumEntry): IrEnumEntrySymbol {
-        return enumEntryCache[enumEntry]
+        enumEntryCache[enumEntry]?.let { return it }
+
+        val symbol = IrEnumEntrySymbolImpl()
+        enumEntryCache[enumEntry] = symbol
+
+        val irParent = declarationStorage.findIrParent(enumEntry, fakeOverrideOwnerLookupTag = null) as IrClass
+        if (irParent.isExternalParent()) {
+            classifiersGenerator.createIrEnumEntry(
+                enumEntry,
+                irParent = irParent,
+                symbol,
+                predefinedOrigin = irParent.origin
+            )
+        }
+        return symbol
     }
 
-    fun getOrCreateIrEnumEntry(
+    fun createAndCacheIrEnumEntry(
         enumEntry: FirEnumEntry,
         irParent: IrClass,
         predefinedOrigin: IrDeclarationOrigin? = null,
     ): IrEnumEntry {
-        getCachedIrEnumEntry(enumEntry)?.let { return it }
+        val symbol = getIrEnumEntrySymbol(enumEntry)
-
         val containingFile = firProvider.getFirCallableContainerFile(enumEntry.symbol)
 
         @Suppress("NAME_SHADOWING")
@@ -351,15 +364,13 @@ class Fir2IrClassifierStorage(
         } else {
             irParent.origin
         }
-        val symbol = IrEnumEntrySymbolImpl()
+
         return classifiersGenerator.createIrEnumEntry(
             enumEntry,
             irParent = irParent,
             symbol,
             predefinedOrigin = predefinedOrigin
-        ).also {
+        )
-            enumEntryCache[enumEntry] = it
-        }
     }
 
     // ------------------------------------ typealiases ------------------------------------

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrCommonMemberStorage.kt

@@ -30,7 +30,7 @@ class Fir2IrCommonMemberStorage(firMangler: FirMangler) {
 
     val typeParameterCache: MutableMap<FirTypeParameter, IrTypeParameter> = mutableMapOf()
 
-    val enumEntryCache: MutableMap<FirEnumEntry, IrEnumEntry> = mutableMapOf()
+    val enumEntryCache: MutableMap<FirEnumEntry, IrEnumEntrySymbol> = mutableMapOf()
 
     val localClassCache: MutableMap<FirClass, IrClass> = mutableMapOf()
 

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt

@@ -462,7 +462,7 @@ class Fir2IrConverter(
     private fun processClassAndNestedClassHeaders(klass: FirClass) {
         classifiersGenerator.processClassHeader(klass)
         processNestedClassHeaders(klass)
-        val irClass = classifierStorage.getIrClass(klass)!!
+        val irClass = classifierStorage.getIrClass(klass)
         /*
          * This is needed to preserve the source order of declarations in the class
          * IrClass should contain declarations in the source order, but creating of nested IrClass automatically adds created class to the list
@@ -521,7 +521,7 @@ class Fir2IrConverter(
         val isInLocalClass = containingClass != null && (containingClass !is FirRegularClass || containingClass.isLocal)
         when (declaration) {
             is FirRegularClass -> {
-                val irClass = classifierStorage.getIrClass(declaration)!!
+                val irClass = classifierStorage.getIrClass(declaration)
                 addDeclarationToParentIfNeeded(irClass)
                 processClassMembers(declaration, irClass)
             }
@@ -594,7 +594,7 @@ class Fir2IrConverter(
                 declarationStorage.createAndCacheIrConstructor(declaration, { parent as IrClass }, isLocal = isInLocalClass)
             }
             is FirEnumEntry -> {
-                classifierStorage.getOrCreateIrEnumEntry(declaration, parent as IrClass)
+                classifierStorage.createAndCacheIrEnumEntry(declaration, parent as IrClass)
             }
             is FirAnonymousInitializer -> {
                 declarationStorage.createIrAnonymousInitializer(declaration, parent as IrClass)

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrDeclarationStorage.kt

@@ -992,26 +992,12 @@ class Fir2IrDeclarationStorage(
 
     fun getIrValueSymbol(firVariableSymbol: FirVariableSymbol<*>): IrSymbol {
         return when (val firDeclaration = firVariableSymbol.fir) {
-            is FirEnumEntry -> {
+            is FirEnumEntry -> classifierStorage.getIrEnumEntrySymbol(firDeclaration)
-                classifierStorage.getCachedIrEnumEntry(firDeclaration)?.let { return it.symbol }
-                val irParentClass = firDeclaration.containingClassLookupTag()?.let { classifierStorage.getIrClass(it) }!!
 
-                val containingFile = firProvider.getFirCallableContainerFile(firVariableSymbol)
+            is FirValueParameter -> localStorage.getParameter(firDeclaration)
+                ?: getIrVariableSymbol(firDeclaration) // catch parameter is FirValueParameter in FIR but IrVariable in IR
 
-                classifierStorage.getOrCreateIrEnumEntry(
+            else -> getIrVariableSymbol(firDeclaration)
-                    firDeclaration,
-                    irParent = irParentClass,
-                    predefinedOrigin = if (containingFile != null) IrDeclarationOrigin.DEFINED else irParentClass.origin
-                ).symbol
-            }
-            is FirValueParameter -> {
-                localStorage.getParameter(firDeclaration)
-                // catch parameter is FirValueParameter in FIR but IrVariable in IR
-                    ?: return getIrVariableSymbol(firDeclaration)
-            }
-            else -> {
-                getIrVariableSymbol(firDeclaration)
-            }
         }
     }
 
@@ -1498,7 +1484,7 @@ class Fir2IrDeclarationStorage(
         return parentPackage
     }
 
-    private fun findIrParent(
+    internal fun findIrParent(
         callableDeclaration: FirCallableDeclaration,
         fakeOverrideOwnerLookupTag: ConeClassLikeLookupTag?,
     ): IrDeclarationParent? {

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -126,7 +126,9 @@ class Fir2IrVisitor(
     }
 
     override fun visitEnumEntry(enumEntry: FirEnumEntry, data: Any?): IrElement = whileAnalysing(session, enumEntry) {
-        val irEnumEntry = classifierStorage.getCachedIrEnumEntry(enumEntry)!!
+        // At this point all IR for source enum entries should be created and bound to symbols
+        @OptIn(UnsafeDuringIrConstructionAPI::class)
+        val irEnumEntry = classifierStorage.getIrEnumEntrySymbol(enumEntry).owner
         annotationGenerator.generate(irEnumEntry, enumEntry)
         val correspondingClass = irEnumEntry.correspondingClass
         val initializer = enumEntry.initializer
@@ -191,7 +193,7 @@ class Fir2IrVisitor(
         if (regularClass.visibility == Visibilities.Local) {
             val irParent = conversionScope.parentFromStack()
             // NB: for implicit types it is possible that local class is already cached
-            val irClass = classifierStorage.getIrClass(regularClass)?.apply { this.parent = irParent }
+            val irClass = classifierStorage.getCachedIrLocalClass(regularClass)?.apply { this.parent = irParent }
             if (irClass != null) {
                 conversionScope.withParent(irClass) {
                     memberGenerator.convertClassContent(irClass, regularClass)
@@ -200,7 +202,7 @@ class Fir2IrVisitor(
             }
             converter.processLocalClassAndNestedClasses(regularClass, irParent)
         }
-        val irClass = classifierStorage.getIrClass(regularClass)!!
+        val irClass = classifierStorage.getIrClass(regularClass)
         if (regularClass.isSealed) {
             irClass.sealedSubclasses = regularClass.getIrSymbolsForSealedSubclasses()
         }
@@ -711,7 +713,7 @@ class Fir2IrVisitor(
             // We anyway can use 'else' branch as fallback, but
             // this is an additional check of FIR2IR invariants
             // (source classes should be already built when we analyze bodies)
-            classifierStorage.getIrClass(firClass)!!.symbol
+            classifierStorage.getIrClass(firClass).symbol
         } else {
             /*
              * The only case when we can refer to non-source this is resolution to companion object of parent

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/FirIrProvider.kt

@@ -172,15 +172,14 @@ class FirIrProvider(val components: Fir2IrComponents) : IrProvider {
         }
 
         val firDeclaration = findDeclarationByHash(firCandidates, signature.id) ?: return null
-        val parent = parentClass ?: declarationStorage.getIrExternalPackageFragment(packageFqName, firDeclaration.moduleData)
 
         return when (kind) {
             SymbolKind.CLASS_SYMBOL -> {
                 shouldNotBeCalled()
             }
-            SymbolKind.ENUM_ENTRY_SYMBOL -> classifierStorage.getOrCreateIrEnumEntry(
+            SymbolKind.ENUM_ENTRY_SYMBOL -> {
-                firDeclaration as FirEnumEntry, parent as IrClass
+                shouldNotBeCalled()
-            )
+            }
             SymbolKind.CONSTRUCTOR_SYMBOL -> {
                 shouldNotBeCalled()
             }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/lazy/Fir2IrLazyClass.kt

@@ -181,7 +181,7 @@ class Fir2IrLazyClass(
         if (fir.classKind == ClassKind.ENUM_CLASS) {
             for (declaration in fir.declarations) {
                 if (declaration is FirEnumEntry && shouldBuildStub(declaration)) {
-                    result += classifierStorage.getOrCreateIrEnumEntry(declaration, this, origin)
+                    result += classifierStorage.getIrEnumEntrySymbol(declaration).owner
                 }
             }
         }

compiler/testData/codegen/box/annotations/selfReferentialAnnotation.fir.ir.txt

@@ -60,10 +60,10 @@ FILE fqName:<root> fileName:/selfReferentialAnnotation.kt
               receiver: GET_VAR '<this>: <root>.MyRequiresOptIn declared in <root>.MyRequiresOptIn.<get-b>' type=<root>.MyRequiresOptIn origin=null
     CLASS ENUM_CLASS name:MyLevel modality:FINAL visibility:public superTypes:[kotlin.Enum<<root>.MyRequiresOptIn.MyLevel>]
       $this: VALUE_PARAMETER INSTANCE_RECEIVER name:<this> type:<root>.MyRequiresOptIn.MyLevel
-      ENUM_ENTRY name:ERROR
+      ENUM_ENTRY name:WARNING
         init: EXPRESSION_BODY
           ENUM_CONSTRUCTOR_CALL 'private constructor <init> () declared in <root>.MyRequiresOptIn.MyLevel'
-      ENUM_ENTRY name:WARNING
+      ENUM_ENTRY name:ERROR
         init: EXPRESSION_BODY
           ENUM_CONSTRUCTOR_CALL 'private constructor <init> () declared in <root>.MyRequiresOptIn.MyLevel'
       CONSTRUCTOR visibility:private <> () returnType:<root>.MyRequiresOptIn.MyLevel [primary]

compiler/testData/codegen/box/annotations/selfReferentialAnnotation.kt

@@ -1,7 +1,6 @@
 // FIR_DUMP
 // DUMP_IR
 // WITH_STDLIB
-// IGNORE_BACKEND_K2: ANY
 
 annotation class Ann(@Ann(1) val e: Int)