[FIR] Change nested class generation extension to generate top-level classes also

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirSessionCommonComponents.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir
 import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedNestedClassIndex
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClassIndex
 import org.jetbrains.kotlin.fir.scopes.impl.FirDeclaredMemberScopeProvider
 import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache
 
@@ -19,5 +19,5 @@ fun FirSession.registerCommonComponents() {
     register(FirExtensionService::class, FirExtensionService(this))
     register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.create(this))
     register(FirPredicateBasedProvider::class, FirPredicateBasedProvider.create(this))
-    register(GeneratedNestedClassIndex::class, GeneratedNestedClassIndex.create())
+    register(GeneratedClassIndex::class, GeneratedClassIndex.create())
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirDeclarationGenerationExtension.kt

@@ -9,14 +9,14 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedNestedClass
+import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClass
 import kotlin.reflect.KClass
 
 /*
  * TODO:
  *  - check that annotations or meta-annotations is not empty
  */
-abstract class FirExistingClassModificationExtension(session: FirSession) : FirPredicateBasedExtension(session) {
+abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredicateBasedExtension(session) {
     companion object {
         val NAME = FirExtensionPointName("ExistingClassModification")
     }
@@ -24,23 +24,23 @@ abstract class FirExistingClassModificationExtension(session: FirSession) : FirP
     final override val name: FirExtensionPointName
         get() = NAME
 
-    final override val extensionType: KClass<out FirExtension> = FirExistingClassModificationExtension::class
+    final override val extensionType: KClass<out FirExtension> = FirDeclarationGenerationExtension::class
 
-    abstract fun generateNestedClasses(
+    abstract fun generateClasses(
         annotatedDeclaration: FirDeclaration,
         owners: List<FirAnnotatedDeclaration>
     ): List<GeneratedDeclaration<FirRegularClass>>
 
-    abstract fun generateMembersForNestedClasses(generatedNestedClass: GeneratedNestedClass): List<FirDeclaration>
+    abstract fun generateMembersForGeneratedClass(generatedClass: GeneratedClass): List<FirDeclaration>
 
     abstract fun generateMembers(
         annotatedDeclaration: FirDeclaration,
         owners: List<FirAnnotatedDeclaration>
     ): List<GeneratedDeclaration<*>>
 
-    data class GeneratedDeclaration<out T : FirDeclaration>(val newDeclaration: T, val owner: FirRegularClass)
+    data class GeneratedDeclaration<out T : FirDeclaration>(val newDeclaration: T, val owner: FirAnnotatedDeclaration)
 
-    fun interface Factory : FirExtension.Factory<FirExistingClassModificationExtension>
+    fun interface Factory : FirExtension.Factory<FirDeclarationGenerationExtension>
 }
 
-val FirExtensionService.existingClassModifiers: List<FirExistingClassModificationExtension> by FirExtensionService.registeredExtensions()
+val FirExtensionService.declarationGenerators: List<FirDeclarationGenerationExtension> by FirExtensionService.registeredExtensions()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionRegistrar.kt

@@ -16,7 +16,7 @@ abstract class FirExtensionRegistrar {
     ) {
         val AVAILABLE_EXTENSIONS = listOf(
             FirStatusTransformerExtension::class,
-            FirExistingClassModificationExtension::class,
+            FirDeclarationGenerationExtension::class,
             AbstractFirAdditionalCheckersExtension::class
         )
     }
@@ -30,8 +30,8 @@ abstract class FirExtensionRegistrar {
         }
 
         @JvmName("plusClassGenerationExtension")
-        operator fun ((FirSession) -> FirExistingClassModificationExtension).unaryPlus() {
-            registerExtension(FirExistingClassModificationExtension::class, FirExistingClassModificationExtension.Factory { this.invoke(it) })
+        operator fun ((FirSession) -> FirDeclarationGenerationExtension).unaryPlus() {
+            registerExtension(FirDeclarationGenerationExtension::class, FirDeclarationGenerationExtension.Factory { this.invoke(it) })
         }
 
         @JvmName("plusAdditionalCheckersExtension")

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt

@@ -25,11 +25,23 @@ abstract class FirPredicateBasedProvider : FirSessionComponent {
     }
 
     abstract fun getSymbolsByPredicate(predicate: DeclarationPredicate): List<FirAnnotatedDeclaration>
+
+    abstract fun getSymbolsByPredicate(
+        declarations: Collection<FirAnnotatedDeclaration>,
+        predicate: DeclarationPredicate
+    ): List<FirAnnotatedDeclaration>
+
     abstract fun getSymbolsWithOwnersByPredicate(
         predicate: DeclarationPredicate
     ): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>>
 
+    abstract fun getSymbolsWithOwnersByPredicate(
+        declarations: Collection<FirAnnotatedDeclaration>,
+        predicate: DeclarationPredicate
+    ): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>>
+
     abstract fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>)
+    abstract fun registerGeneratedDeclaration(declaration: FirAnnotatedDeclaration, owner: FirAnnotatedDeclaration)
 
     abstract fun matches(predicate: DeclarationPredicate, declaration: FirAnnotatedDeclaration): Boolean
 }
@@ -41,13 +53,30 @@ private class FirPredicateBasedProviderImpl(private val session: FirSession) : F
     override fun getSymbolsByPredicate(predicate: DeclarationPredicate): List<FirAnnotatedDeclaration> {
         val annotations = registeredPluginAnnotations.getAnnotationsForPredicate(predicate)
         if (annotations.isEmpty()) return emptyList()
-        return annotations.flatMapTo(mutableSetOf()) { cache.declarationByAnnotation[it] + cache.declarationsUnderAnnotated[it] }.filter {
-            matches(predicate, it)
-        }
+        val declarations = annotations.flatMapTo(mutableSetOf()) { cache.declarationByAnnotation[it] + cache.declarationsUnderAnnotated[it] }
+        return getSymbolsByPredicate(declarations, predicate)
+    }
+
+    override fun getSymbolsByPredicate(
+        declarations: Collection<FirAnnotatedDeclaration>,
+        predicate: DeclarationPredicate
+    ): List<FirAnnotatedDeclaration> {
+        return declarations.filter { matches(predicate, it) }
     }
 
     override fun getSymbolsWithOwnersByPredicate(predicate: DeclarationPredicate): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>> {
-        return getSymbolsByPredicate(predicate).map { it to cache.ownersForDeclaration.getValue(it) }
+        return getSymbolsByPredicate(predicate).zipWithParents()
+    }
+
+    override fun getSymbolsWithOwnersByPredicate(
+        declarations: Collection<FirAnnotatedDeclaration>,
+        predicate: DeclarationPredicate
+    ): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>> {
+        return getSymbolsByPredicate(declarations, predicate).zipWithParents()
+    }
+
+    private fun List<FirAnnotatedDeclaration>.zipWithParents(): List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>> {
+        return this.map { it to cache.ownersForDeclaration.getValue(it) }
     }
 
     override fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>) {
@@ -62,6 +91,11 @@ private class FirPredicateBasedProviderImpl(private val session: FirSession) : F
         cache.annotationsOfDeclaration.putAll(declaration, matchingAnnotations)
     }
 
+    override fun registerGeneratedDeclaration(declaration: FirAnnotatedDeclaration, owner: FirAnnotatedDeclaration) {
+        val owners = cache.ownersForDeclaration.getValue(owner).add(owner)
+        registerAnnotatedDeclaration(declaration, owners)
+    }
+
     private fun registerOwnersDeclarations(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>) {
         val lastOwner = owners.lastOrNull() ?: return
         val annotationsFromLastOwner = cache.annotationsOfDeclaration[lastOwner]

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/FirProvider.kt

@@ -5,9 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.providers
 
-import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.name.ClassId
@@ -47,5 +45,8 @@ abstract class FirProvider : FirSymbolProvider() {
     abstract fun getFirFilesByPackage(fqName: FqName): List<FirFile>
 
     @FirProviderInternals
-    abstract fun recordNestedClass(owner: FirRegularClass, klass: FirRegularClass)
+    abstract fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass)
+
+    @FirProviderInternals
+    abstract fun recordGeneratedMember(owner: FirAnnotatedDeclaration, klass: FirDeclaration)
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirProviderImpl.kt

@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
-import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -68,11 +67,22 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: KotlinSc
     }
 
     @FirProviderInternals
-    override fun recordNestedClass(owner: FirRegularClass, klass: FirRegularClass) {
-        val file = getFirClassifierContainerFile(owner.symbol.classId)
-        klass.accept(FirRecorder, state to file)
+    override fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass) {
+        klass.accept(FirRecorder, state to owner.file)
     }
 
+    @FirProviderInternals
+    override fun recordGeneratedMember(owner: FirAnnotatedDeclaration, klass: FirDeclaration) {
+        klass.accept(FirRecorder, state to owner.file)
+    }
+
+    private val FirAnnotatedDeclaration.file: FirFile
+        get() = when (this) {
+            is FirFile -> this
+            is FirRegularClass -> getFirClassifierContainerFile(this.symbol.classId)
+            else -> error("Should not be here")
+        }
+
     private fun recordFile(file: FirFile, state: State) {
         val packageName = file.packageFqName
         state.fileMap.merge(packageName, listOf(file)) { a, b -> a + b }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/ResolvePhaseUtils.kt

@@ -25,7 +25,7 @@ fun FirResolvePhase.createCompilerProcessorByPhase(
     return when (this) {
         RAW_FIR -> throw IllegalStateException("Raw FIR building phase does not have a transformer")
         ANNOTATIONS_FOR_PLUGINS -> FirPluginAnnotationsResolveProcessor(session, scopeSession)
-        NESTED_CLASS_GENERATION -> FirGlobalNestedClassGenerationProcessor(session, scopeSession)
+        CLASS_GENERATION -> FirGlobalClassGenerationProcessor(session, scopeSession)
         IMPORTS -> FirImportResolveProcessor(session, scopeSession)
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)
@@ -46,7 +46,7 @@ fun FirResolvePhase.createTransformerBasedProcessorByPhase(
     return when (this) {
         RAW_FIR -> throw IllegalStateException("Raw FIR building phase does not have a transformer")
         ANNOTATIONS_FOR_PLUGINS -> FirPluginAnnotationsResolveProcessor(session, scopeSession)
-        NESTED_CLASS_GENERATION -> FirDummyTransformerBasedProcessor(session, scopeSession) // TODO: remove
+        CLASS_GENERATION -> FirDummyTransformerBasedProcessor(session, scopeSession) // TODO: remove
         IMPORTS -> FirImportResolveProcessor(session, scopeSession)
         SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
         SEALED_CLASS_INHERITORS -> FirSealedClassInheritorsProcessor(session, scopeSession)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalClassGenerationProcessor.kt

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2010-2020 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.fir.resolve.transformers.plugin
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.addDeclaration
+import org.jetbrains.kotlin.fir.extensions.*
+import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.firProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
+import org.jetbrains.kotlin.fir.resolve.transformers.FirGlobalResolveProcessor
+
+class FirGlobalClassGenerationProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirGlobalResolveProcessor(session, scopeSession) {
+    override fun process() {
+        val extensions = session.extensionService.declarationGenerators
+        if (extensions.isEmpty()) return
+        val provider = session.predicateBasedProvider
+        for (extension in extensions) {
+            var annotatedDeclarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
+            while (annotatedDeclarations.isNotEmpty()) {
+                val newClasses = generateClasses(annotatedDeclarations, extension)
+                annotatedDeclarations = provider.getSymbolsWithOwnersByPredicate(newClasses, extension.predicate)
+            }
+        }
+    }
+
+    private fun generateClasses(
+        declarations: List<Pair<FirAnnotatedDeclaration, List<FirAnnotatedDeclaration>>>,
+        extension: FirDeclarationGenerationExtension,
+    ): List<FirRegularClass> {
+        val newClasses = mutableListOf<FirRegularClass>()
+        for ((declaration, owners) in declarations) {
+            generateClass(extension, declaration, owners, newClasses)
+        }
+        return newClasses
+    }
+
+    @OptIn(FirProviderInternals::class)
+    private fun generateClass(
+        extension: FirDeclarationGenerationExtension,
+        declaration: FirAnnotatedDeclaration,
+        owners: List<FirAnnotatedDeclaration>,
+        newClasses: MutableList<FirRegularClass>
+    ) {
+        val generatedClasses = extension.generateClasses(declaration, owners)
+        for ((klass, owner) in generatedClasses) {
+            when (owner) {
+                is FirRegularClass -> owner.addDeclaration(klass)
+                is FirFile -> owner.addDeclaration(klass)
+            }
+            session.generatedClassIndex.registerClass(klass, owner)
+            session.predicateBasedProvider.registerGeneratedDeclaration(klass, owner)
+            session.firProvider.recordGeneratedClass(owner, klass)
+            newClasses += klass
+        }
+    }
+}
+

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalNestedClassGenerationProcessor.kt

@@ -1,41 +0,0 @@
-/*
- * Copyright 2010-2020 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.fir.resolve.transformers.plugin
-
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.addDeclaration
-import org.jetbrains.kotlin.fir.extensions.existingClassModifiers
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.firProvider
-import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
-import org.jetbrains.kotlin.fir.resolve.transformers.FirGlobalResolveProcessor
-
-class FirGlobalNestedClassGenerationProcessor(
-    session: FirSession,
-    scopeSession: ScopeSession
-) : FirGlobalResolveProcessor(session, scopeSession) {
-    @OptIn(FirProviderInternals::class)
-    override fun process() {
-        val extensions = session.extensionService.existingClassModifiers
-        if (extensions.isEmpty()) return
-        val provider = session.predicateBasedProvider
-        val index = session.generatedNestedClassIndex
-        for (extension in extensions) {
-            val declarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
-            for ((declaration, owners) in declarations) {
-                val nestedClasses = extension.generateNestedClasses(declaration, owners)
-                for ((nestedClass, owner) in nestedClasses) {
-                    owner.addDeclaration(nestedClass)
-                    index.registerNestedClass(nestedClass, owner)
-                    session.firProvider.recordNestedClass(owner, nestedClass)
-                }
-            }
-        }
-    }
-}
-

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirGlobalNewMemberGenerationProcessor.kt

@@ -6,49 +6,61 @@
 package org.jetbrains.kotlin.fir.resolve.transformers.plugin
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.addDeclaration
 import org.jetbrains.kotlin.fir.declarations.validate
-import org.jetbrains.kotlin.fir.extensions.FirExistingClassModificationExtension
-import org.jetbrains.kotlin.fir.extensions.existingClassModifiers
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
+import org.jetbrains.kotlin.fir.extensions.declarationGenerators
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.firProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
 import org.jetbrains.kotlin.fir.resolve.transformers.FirGlobalResolveProcessor
 
 class FirGlobalNewMemberGenerationProcessor(
     session: FirSession,
     scopeSession: ScopeSession
 ) : FirGlobalResolveProcessor(session, scopeSession) {
-    private val index = session.generatedNestedClassIndex
+    private val index = session.generatedClassIndex
     private val provider = session.predicateBasedProvider
 
     override fun process() {
-        val extensions = session.extensionService.existingClassModifiers
+        val extensions = session.extensionService.declarationGenerators
         if (extensions.isEmpty()) return
         for (extension in extensions) {
             generateNewMembers(extension)
-            fillGeneratedNestedClasses(extension)
+            fillGeneratedClasses(extension)
         }
     }
 
-    private fun generateNewMembers(extension: FirExistingClassModificationExtension) {
+    @OptIn(FirProviderInternals::class)
+    private fun generateNewMembers(extension: FirDeclarationGenerationExtension) {
         val declarations = provider.getSymbolsWithOwnersByPredicate(extension.predicate)
         for ((declaration, owners) in declarations) {
             val newMembers = extension.generateMembers(declaration, owners)
             for ((newMember, owner) in newMembers) {
                 newMember.validate()
-                owner.addDeclaration(newMember)
+                when (owner) {
+                    is FirRegularClass -> owner.addDeclaration(newMember)
+                    is FirFile -> owner.addDeclaration(newMember)
+                    else -> error("Should not be here")
+                }
+                session.firProvider.recordGeneratedMember(owner, newMember)
             }
         }
     }
 
-    private fun fillGeneratedNestedClasses(extension: FirExistingClassModificationExtension) {
-        for (generatedNestedClass in index[extension.key]) {
-            val nestedClass = generatedNestedClass.nestedClass
-            val newMembers = extension.generateMembersForNestedClasses(generatedNestedClass)
+    @OptIn(FirProviderInternals::class)
+    private fun fillGeneratedClasses(extension: FirDeclarationGenerationExtension) {
+        for (generatedClass in index[extension.key]) {
+            val klass = generatedClass.klass
+            val newMembers = extension.generateMembersForGeneratedClass(generatedClass)
             for (newMember in newMembers) {
                 newMember.validate()
-                nestedClass.addDeclaration(newMember)
+                klass.addDeclaration(newMember)
+                session.firProvider.recordGeneratedMember(klass, newMember)
             }
         }
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/GeneratedClassIndex.kt

@@ -8,34 +8,35 @@ package org.jetbrains.kotlin.fir.resolve.transformers.plugin
 import com.google.common.collect.ArrayListMultimap
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 
-data class GeneratedNestedClass(val nestedClass: FirRegularClass, val owner: FirRegularClass)
+data class GeneratedClass(val klass: FirRegularClass, val owner: FirAnnotatedDeclaration)
 
-abstract class GeneratedNestedClassIndex : FirSessionComponent {
+abstract class GeneratedClassIndex : FirSessionComponent {
     companion object {
-        fun create(): GeneratedNestedClassIndex {
-            return GeneratedNestedClassIndexImpl()
+        fun create(): GeneratedClassIndex {
+            return GeneratedClassIndexImpl()
         }
     }
 
-    abstract fun registerNestedClass(klass: FirRegularClass, owner: FirRegularClass)
-    abstract operator fun get(key: FirPluginKey): List<GeneratedNestedClass>
+    abstract fun registerClass(klass: FirRegularClass, owner: FirAnnotatedDeclaration)
+    abstract operator fun get(key: FirPluginKey): List<GeneratedClass>
 }
 
-val FirSession.generatedNestedClassIndex: GeneratedNestedClassIndex by FirSession.sessionComponentAccessor()
+val FirSession.generatedClassIndex: GeneratedClassIndex by FirSession.sessionComponentAccessor()
 
-private class GeneratedNestedClassIndexImpl : GeneratedNestedClassIndex() {
-    private val index: ArrayListMultimap<FirPluginKey, GeneratedNestedClass> = ArrayListMultimap.create()
+private class GeneratedClassIndexImpl : GeneratedClassIndex() {
+    private val index: ArrayListMultimap<FirPluginKey, GeneratedClass> = ArrayListMultimap.create()
 
-    override fun registerNestedClass(klass: FirRegularClass, owner: FirRegularClass) {
+    override fun registerClass(klass: FirRegularClass, owner: FirAnnotatedDeclaration) {
         val key = (klass.origin as FirDeclarationOrigin.Plugin).key
-        index.put(key, GeneratedNestedClass(klass, owner))
+        index.put(key, GeneratedClass(klass, owner))
     }
 
-    override fun get(key: FirPluginKey): List<GeneratedNestedClass> {
+    override fun get(key: FirPluginKey): List<GeneratedClass> {
         return index.get(key)
     }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirResolvePhase.kt

@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.fir.declarations
 enum class FirResolvePhase {
     RAW_FIR,
     ANNOTATIONS_FOR_PLUGINS, // run only if some extensions are registered
-    NESTED_CLASS_GENERATION, // plugin phase
+    CLASS_GENERATION, // plugin phase
     IMPORTS,
     SUPER_TYPES,
     SEALED_CLASS_INHERITORS,

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeProvider.kt

@@ -10,9 +10,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirProviderInternals
@@ -153,9 +151,15 @@ class FirIdeProvider(
     }
 
     @FirProviderInternals
-    override fun recordNestedClass(owner: FirRegularClass, klass: FirRegularClass) {
+    override fun recordGeneratedClass(owner: FirAnnotatedDeclaration, klass: FirRegularClass) {
         // TODO: check that this implementation is correct
-        cacheProvider.recordNestedClass(owner, klass)
+        cacheProvider.recordGeneratedClass(owner, klass)
+    }
+
+    @FirProviderInternals
+    override fun recordGeneratedMember(owner: FirAnnotatedDeclaration, klass: FirDeclaration) {
+        // TODO: check that this implementation is correct
+        cacheProvider.recordGeneratedMember(owner, klass)
     }
 }