FIR2IR: split local class on-the-fly creation into two phases

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

@@ -44,6 +44,10 @@ class Fir2IrClassifierStorage(
 
     private val classCache = mutableMapOf<FirRegularClass, IrClass>()
 
+    private val localClassesCreatedOnTheFly = mutableMapOf<FirClass, IrClass>()
+
+    private var processMembersOfClassesOnTheFlyImmediately = false
+
     private val typeAliasCache = mutableMapOf<FirTypeAlias, IrTypeAlias>()
 
     private val typeParameterCache = mutableMapOf<FirTypeParameter, IrTypeParameter>()
@@ -168,16 +172,41 @@ class Fir2IrClassifierStorage(
     // This function is called when we refer local class earlier than we reach its declaration
     // This can happen e.g. when implicit return type has a local class constructor
     private fun createLocalIrClassOnTheFly(klass: FirClass): IrClass {
-        // NB: klass can be either FirRegularClass or FirAnonymousObject
         val result = when (klass) {
             is FirAnonymousObject -> createIrAnonymousObject(klass, irParent = temporaryParent).apply {
-                converter.processAnonymousObjectMembers(klass, this)
+                converter.processAnonymousObjectOnTheFly(klass, this)
             }
-            is FirRegularClass -> converter.processLocalClassAndNestedClasses(klass, temporaryParent)
+            is FirRegularClass -> converter.processLocalClassAndNestedClassesOnTheFly(klass, temporaryParent)
+        }
+        // Note: usually member creation and f/o binding is delayed till non-local classes are processed in Fir2IrConverter
+        // If non-local classes are already created (this means we are in body translation) we do everything immediately
+        // The last variant is possible for local variables like 'val a = object : Any() { ... }'
+        if (processMembersOfClassesOnTheFlyImmediately) {
+            processMembersOfClassCreatedOnTheFly(klass, result)
+            converter.bindFakeOverridesInClass(result)
+        } else {
+            localClassesCreatedOnTheFly[klass] = result
         }
         return result
     }
 
+    fun processMembersOfClassesCreatedOnTheFly() {
+        processMembersOfClassesOnTheFlyImmediately = true
+        for ((klass, irClass) in localClassesCreatedOnTheFly) {
+            processMembersOfClassCreatedOnTheFly(klass, irClass)
+        }
+        // Note: it's better to bind everything AFTER members are built, in case local classes are dependent on each other
+        localClassesCreatedOnTheFly.values.forEach(converter::bindFakeOverridesInClass)
+        localClassesCreatedOnTheFly.clear()
+    }
+
+    private fun processMembersOfClassCreatedOnTheFly(klass: FirClass, irClass: IrClass) {
+        when (klass) {
+            is FirRegularClass -> converter.processClassMembers(klass, irClass)
+            is FirAnonymousObject -> converter.processAnonymousObjectMembers(klass, irClass, processHeaders = false)
+        }
+    }
+
     fun processClassHeader(regularClass: FirRegularClass, irClass: IrClass = getCachedIrClass(regularClass)!!): IrClass {
         irClass.declareTypeParameters(regularClass)
         irClass.setThisReceiver(regularClass.typeParameters)

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

@@ -52,6 +52,17 @@ class Fir2IrConverter(
 
     private val generatorExtensions = session.extensionService.declarationGenerators
 
+    fun processLocalClassAndNestedClassesOnTheFly(regularClass: FirRegularClass, parent: IrDeclarationParent): IrClass {
+        val irClass = registerClassAndNestedClasses(regularClass, parent)
+        processClassAndNestedClassHeaders(regularClass)
+        return irClass
+    }
+
+    fun processAnonymousObjectOnTheFly(anonymousObject: FirAnonymousObject, irClass: IrClass) {
+        registerNestedClasses(anonymousObject, irClass)
+        processNestedClassHeaders(anonymousObject)
+    }
+
     fun processLocalClassAndNestedClasses(regularClass: FirRegularClass, parent: IrDeclarationParent): IrClass {
         val irClass = registerClassAndNestedClasses(regularClass, parent)
         processClassAndNestedClassHeaders(regularClass)
@@ -117,10 +128,13 @@ class Fir2IrConverter(
 
     fun processAnonymousObjectMembers(
         anonymousObject: FirAnonymousObject,
-        irClass: IrClass = classifierStorage.getCachedIrClass(anonymousObject)!!
+        irClass: IrClass,
+        processHeaders: Boolean
     ): IrClass {
-        registerNestedClasses(anonymousObject, irClass)
-        processNestedClassHeaders(anonymousObject)
+        if (processHeaders) {
+            registerNestedClasses(anonymousObject, irClass)
+            processNestedClassHeaders(anonymousObject)
+        }
         anonymousObject.primaryConstructorIfAny(session)?.let {
             irClass.declarations += declarationStorage.createIrConstructor(
                 it.fir, irClass, isLocal = true
@@ -388,6 +402,7 @@ class Fir2IrConverter(
             for (firFile in allFirFiles) {
                 converter.bindFakeOverridesInFile(firFile)
             }
+            classifierStorage.processMembersOfClassesCreatedOnTheFly()
 
             for (firFile in allFirFiles) {
                 firFile.accept(fir2irVisitor, null)

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

@@ -1167,11 +1167,7 @@ class Fir2IrDeclarationStorage(
                 ) as IrFunctionSymbol
                 if (unmatchedReceiver && dispatchReceiverLookupTag is ConeClassLookupTagWithFixedSymbol) {
                     val originalFunction = originalSymbol.owner as IrSimpleFunction
-                    dispatchReceiverLookupTag.findIrFakeOverride(
-                        fir.name, originalFunction
-                    ) as? IrFunctionSymbol
-                    // Fallback for simple synthetic property
-                        ?: originalSymbol
+                    dispatchReceiverLookupTag.findIrFakeOverride(fir.name, originalFunction) as IrFunctionSymbol
                 } else {
                     originalSymbol
                 }

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

@@ -115,7 +115,7 @@ class Fir2IrVisitor(
             declarationStorage.enterScope(irEnumEntry)
             classifierStorage.putEnumEntryClassInScope(enumEntry, correspondingClass)
             val anonymousObject = (enumEntry.initializer as FirAnonymousObjectExpression).anonymousObject
-            converter.processAnonymousObjectMembers(anonymousObject, correspondingClass)
+            converter.processAnonymousObjectMembers(anonymousObject, correspondingClass, processHeaders = true)
             conversionScope.withParent(correspondingClass) {
                 conversionScope.withContainingFirClass(anonymousObject) {
                     memberGenerator.convertClassContent(correspondingClass, anonymousObject)
@@ -177,8 +177,8 @@ class Fir2IrVisitor(
         val irParent = conversionScope.parentFromStack()
         // NB: for implicit types it is possible that anonymous object is already cached
         val irAnonymousObject = classifierStorage.getCachedIrClass(anonymousObject)?.apply { this.parent = irParent }
-            ?: classifierStorage.createIrAnonymousObject(anonymousObject, irParent = irParent).also {
-                converter.processAnonymousObjectMembers(anonymousObject, it)
+            ?: classifierStorage.createIrAnonymousObject(anonymousObject, irParent = irParent).also { irClass ->
+                converter.processAnonymousObjectMembers(anonymousObject, irClass, processHeaders = true)
             }
         conversionScope.withParent(irAnonymousObject) {
             conversionScope.withContainingFirClass(anonymousObject) {