[FIR] Change scheme of generating declarations by plugins

Methods `needToGenerateAdditionalMembersInClass` and `needToGenerateNestedClassifiersInClass` are removed, now compiler uses `get...Names` and `getTopLevel...` methods to determine which extension may generate declaration with specific classId/callableId This is needed to simplify API of FirDeclarationGenerationExtension and provide guarantee that `generate...` method will be called with specific classId/callableId only if specific extensions returned name for this id from `getName...` functions
2021-11-15 12:57:26 +03:00
parent 40d8451698
commit cb0705ec03
25 changed files with 388 additions and 228 deletions
@@ -19,10 +19,7 @@ import org.jetbrains.kotlin.fir.checkers.registerJvmCheckers
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.deserialization.ModuleDataProvider
 import org.jetbrains.kotlin.fir.deserialization.SingleModuleDataProvider
-import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
-import org.jetbrains.kotlin.fir.extensions.FirExtensionDeclarationsSymbolProvider
-import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.registerExtensions
+import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.java.FirCliSession
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
@@ -167,7 +164,7 @@ object FirSessionFactory {
            }.configure()

            val dependenciesSymbolProvider = FirDependenciesSymbolProviderImpl(this)
-            val generatedSymbolsProvider = FirExtensionDeclarationsSymbolProvider.create(this)
+            val generatedSymbolsProvider = FirSwitchableExtensionDeclarationsSymbolProvider.create(this)
            register(
                FirSymbolProvider::class,
                FirCompositeSymbolProvider(
@@ -182,7 +179,7 @@ object FirSessionFactory {
                )
            )

-            generatedSymbolsProvider?.let { register(FirExtensionDeclarationsSymbolProvider::class, it) }
+            generatedSymbolsProvider?.let { register(FirSwitchableExtensionDeclarationsSymbolProvider::class, it) }

            register(
                FirDependenciesSymbolProvider::class,
@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.backend.generators.FakeOverrideGenerator
 import org.jetbrains.kotlin.fir.builder.buildPackageDirective
+import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.builder.buildFile
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
@@ -23,8 +24,10 @@ import org.jetbrains.kotlin.fir.declarations.utils.isInline
 import org.jetbrains.kotlin.fir.declarations.utils.isJava
 import org.jetbrains.kotlin.fir.declarations.utils.visibility
 import org.jetbrains.kotlin.fir.expressions.*
+import org.jetbrains.kotlin.fir.extensions.FirExtensionApiInternals
 import org.jetbrains.kotlin.fir.extensions.declarationGenerators
 import org.jetbrains.kotlin.fir.extensions.extensionService
+import org.jetbrains.kotlin.fir.extensions.generatedDeclarationsSymbolProvider
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.references.FirReference
@@ -610,11 +613,12 @@ fun FirRegularClass.getIrSymbolsForSealedSubclasses(components: Fir2IrComponents
    }.filterIsInstance<IrClassSymbol>()
 }

+@OptIn(FirExtensionApiInternals::class)
 fun FirSession.createFilesWithGeneratedDeclarations(): List<FirFile> {
-    val symbolProvider = symbolProvider
+    val symbolProvider = generatedDeclarationsSymbolProvider ?: return emptyList()
    val declarationGenerators = extensionService.declarationGenerators
-    val topLevelClasses = declarationGenerators.flatMap { it.getTopLevelClassIds() }.groupBy { it.packageFqName }
-    val topLevelCallables = declarationGenerators.flatMap { it.getTopLevelCallableIds() }.groupBy { it.packageName }
+    val topLevelClasses = declarationGenerators.flatMap { it.topLevelClassIdsCache.getValue() }.groupBy { it.packageFqName }
+    val topLevelCallables = declarationGenerators.flatMap { it.topLevelCallableIdsCache.getValue() }.groupBy { it.packageName }

    return buildList {
        for (packageFqName in (topLevelClasses.keys + topLevelCallables.keys)) {
@@ -6,6 +6,9 @@
 package org.jetbrains.kotlin.fir.extensions

 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.FirLazyValue
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
@@ -28,9 +31,6 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredi

    final override val extensionType: KClass<out FirExtension> = FirDeclarationGenerationExtension::class

-    abstract fun needToGenerateAdditionalMembersInClass(klass: FirClass): Boolean
-    abstract fun needToGenerateNestedClassifiersInClass(klass: FirClass): Boolean
-
    /*
     * Can be called on SUPERTYPES stage
     *
@@ -41,18 +41,43 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirPredi
    // Can be called on STATUS stage
    open fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> = emptyList()
    open fun generateProperties(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirPropertySymbol> = emptyList()
-    open fun generateConstructors(callableId: CallableId): List<FirConstructorSymbol> = emptyList()
+    open fun generateConstructors(owner: FirClassSymbol<*>): List<FirConstructorSymbol> = emptyList()

    // Can be called on IMPORTS stage
    open fun hasPackage(packageFqName: FqName): Boolean = false

-    // Can be called after BODY_RESOLVE stage (checkers and fir2ir)
+    /*
+     * Can be called after SUPERTYPES stage
+     *
+     * `generate...` methods will be called only if `get...Names/ClassIds/CallableIds` returned corresponding
+     *   declaration name
+     *
+     * If you want to generate constructor for some class, then you need to return `SpecialNames.INIT` in
+     *   set of callable names for this class
+     */
    open fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
    open fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> = emptySet()
    open fun getTopLevelCallableIds(): Set<CallableId> = emptySet()
    open fun getTopLevelClassIds(): Set<ClassId> = emptySet()

    fun interface Factory : FirExtension.Factory<FirDeclarationGenerationExtension>
+
+    // ----------------------------------- internal utils -----------------------------------
+
+    @FirExtensionApiInternals
+    val nestedClassifierNamesCache: FirCache<FirClassSymbol<*>, Set<Name>, Nothing?> =
+        session.firCachesFactory.createCache { symbol, _ ->
+            getNestedClassifiersNames(symbol)
+        }
+
+    @FirExtensionApiInternals
+    val topLevelClassIdsCache: FirLazyValue<Set<ClassId>, Nothing?> =
+        session.firCachesFactory.createLazyValue { getTopLevelClassIds() }
+
+    @FirExtensionApiInternals
+    val topLevelCallableIdsCache: FirLazyValue<Set<CallableId>, Nothing?> =
+        session.firCachesFactory.createLazyValue { getTopLevelCallableIds() }
+
 }

 val FirExtensionService.declarationGenerators: List<FirDeclarationGenerationExtension> by FirExtensionService.registeredExtensions()
@@ -7,21 +7,21 @@ package org.jetbrains.kotlin.fir.extensions

 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
-import org.jetbrains.kotlin.fir.caches.FirCache
-import org.jetbrains.kotlin.fir.caches.FirCachesFactory
-import org.jetbrains.kotlin.fir.caches.firCachesFactory
-import org.jetbrains.kotlin.fir.caches.getValue
+import org.jetbrains.kotlin.fir.caches.*
 import org.jetbrains.kotlin.fir.declarations.validate
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.scopes.impl.groupExtensionsByName
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addToStdlib.flatGroupBy

+@OptIn(FirExtensionApiInternals::class)
 class FirExtensionDeclarationsSymbolProvider private constructor(
    session: FirSession,
    cachesFactory: FirCachesFactory,
@@ -53,10 +53,37 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
        hasPackage(packageFqName)
    }

+    private val extensionsByTopLevelClassId: FirLazyValue<Map<ClassId, List<FirDeclarationGenerationExtension>>, Nothing?> =
+        session.firCachesFactory.createLazyValue {
+            extensions.flatGroupBy { it.topLevelClassIdsCache.getValue() }
+        }
+
+    private val extensionsByNestedClassifierClassId: FirCache<ClassId, Map<ClassId, List<FirDeclarationGenerationExtension>>, Nothing?> =
+        session.firCachesFactory.createCache cache@{ outerClassId, _ ->
+            val outerClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(outerClassId) as? FirClassSymbol<*>
+                ?: return@cache emptyMap()
+            session.groupExtensionsByName(
+                outerClassSymbol.fir,
+                nameExtractor = { nestedClassifierNamesCache.getValue(outerClassSymbol) },
+                nameTransformer = { outerClassId.createNestedClassId(it) }
+            )
+        }
+
+    private val extensionsByTopLevelCallableId: FirLazyValue<Map<CallableId, List<FirDeclarationGenerationExtension>>, Nothing?> =
+        session.firCachesFactory.createLazyValue {
+            extensions.flatGroupBy { it.topLevelCallableIdsCache.getValue() }
+        }
+
    // ------------------------------------------ generators ------------------------------------------

    private fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? {
-        val generatedClasses = extensions.mapNotNull { it.generateClassLikeDeclaration(classId) }.onEach { it.fir.validate() }
+        val matchedExtensions = when {
+            classId.isNestedClass -> extensionsByNestedClassifierClassId.getValue(classId.outerClassId!!)[classId]
+            else -> extensionsByTopLevelClassId.getValue()[classId]
+        } ?: return null
+        val generatedClasses = matchedExtensions
+            .mapNotNull { it.generateClassLikeDeclaration(classId) }
+            .onEach { it.fir.validate() }
        return when (generatedClasses.size) {
            0 -> null
            1 -> generatedClasses.first()
@@ -65,11 +92,15 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
    }

    private fun generateTopLevelFunctions(callableId: CallableId): List<FirNamedFunctionSymbol> {
-        return extensions.flatMap { it.generateFunctions(callableId, owner = null) }.onEach { it.fir.validate() }
+        return extensionsByTopLevelCallableId.getValue()[callableId].orEmpty()
+            .flatMap { it.generateFunctions(callableId, owner = null) }
+            .onEach { it.fir.validate() }
    }

    private fun generateTopLevelProperties(callableId: CallableId): List<FirPropertySymbol> {
-        return extensions.flatMap { it.generateProperties(callableId, owner = null) }.onEach { it.fir.validate() }
+        return extensionsByTopLevelCallableId.getValue()[callableId].orEmpty()
+            .flatMap { it.generateProperties(callableId, owner = null) }
+            .onEach { it.fir.validate() }
    }

    private fun hasPackage(packageFqName: FqName): Boolean {
@@ -103,5 +134,3 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
        return fqName.takeIf { packageCache.getValue(fqName, null) }
    }
 }
-
-val FirSession.generatedDeclarationsSymbolProvider: FirExtensionDeclarationsSymbolProvider? by FirSession.nullableSessionComponentAccessor()
@@ -0,0 +1,76 @@
+/*
+ * Copyright 2010-2021 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.extensions
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+/*
+ * This provider is needed because we need to have ability to disable FirExtensionDeclarationsSymbolProvider during
+ *   phase of annotations for plugins resolution. At this stage predicateBasedProvider is not indexed, so it will return
+ *   empty results for all requests
+ *
+ * This is also legal, because plugins can not generate annotation classes which can influence other plugins or this plugin itself
+ */
+class FirSwitchableExtensionDeclarationsSymbolProvider private constructor(
+    private val delegate: FirSymbolProvider
+) : FirSymbolProvider(delegate.session) {
+    companion object {
+        fun create(session: FirSession): FirSwitchableExtensionDeclarationsSymbolProvider? {
+            return FirExtensionDeclarationsSymbolProvider.create(session)?.let { FirSwitchableExtensionDeclarationsSymbolProvider(it) }
+        }
+    }
+
+    private var disabled: Boolean = false
+
+    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
+        if (disabled) return null
+        return delegate.getClassLikeSymbolByClassId(classId)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
+        if (disabled) return
+        delegate.getTopLevelCallableSymbolsTo(destination, packageFqName, name)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
+        if (disabled) return
+        delegate.getTopLevelFunctionSymbolsTo(destination, packageFqName, name)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
+        if (disabled) return
+        delegate.getTopLevelPropertySymbolsTo(destination, packageFqName, name)
+    }
+
+    override fun getPackage(fqName: FqName): FqName? {
+        if (disabled) return null
+        return delegate.getPackage(fqName)
+    }
+
+    @FirSymbolProviderInternals
+    fun disable() {
+        disabled = true
+    }
+
+    @FirSymbolProviderInternals
+    fun enable() {
+        disabled = false
+    }
+}
+
+val FirSession.generatedDeclarationsSymbolProvider: FirSwitchableExtensionDeclarationsSymbolProvider? by FirSession.nullableSessionComponentAccessor()
@@ -12,11 +12,10 @@ import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.FirClass
-import org.jetbrains.kotlin.fir.extensions.declarationGenerators
-import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 import org.jetbrains.kotlin.fir.scopes.FirNameAwareCompositeScope
+import org.jetbrains.kotlin.fir.scopes.FirTypeScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
@@ -31,8 +30,6 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
    private val nestedClassifierCache: FirCache<FirClass, FirNestedClassifierScope?, Nothing?> =
        useSiteSession.firCachesFactory.createCache { klass, _ -> createNestedClassifierScope(klass) }

-    private val extensions by lazy(LazyThreadSafetyMode.PUBLICATION) { useSiteSession.extensionService.declarationGenerators }
-
    fun declaredMemberScope(
        klass: FirClass,
        useLazyNestedClassifierScope: Boolean,
@@ -56,7 +53,7 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
    ): FirContainingNamesAwareScope {
        return when {
            klass.origin.generated -> {
-                FirGeneratedClassDeclaredMemberScope(useSiteSession, klass, needNestedClassifierScope = true)
+                FirGeneratedClassDeclaredMemberScope.create(useSiteSession, klass, needNestedClassifierScope = true) ?: FirTypeScope.Empty
            }
            else -> {
                val baseScope = FirClassDeclaredMemberScopeImpl(
@@ -66,10 +63,9 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
                    existingNames,
                    symbolProvider
                )
-                if (extensions.any { it.needToGenerateAdditionalMembersInClass(klass) }) {
-                    FirNameAwareCompositeScope(
-                        listOf(baseScope, FirGeneratedClassDeclaredMemberScope(useSiteSession, klass, needNestedClassifierScope = false))
-                    )
+                val generatedScope = FirGeneratedClassDeclaredMemberScope.create(useSiteSession, klass, needNestedClassifierScope = false)
+                if (generatedScope != null) {
+                    FirNameAwareCompositeScope(listOf(baseScope, generatedScope))
                } else {
                    baseScope
                }
@@ -83,19 +79,20 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio

    private fun createNestedClassifierScope(klass: FirClass): FirNestedClassifierScope? {
        return if (klass.origin.generated) {
-            FirGeneratedClassNestedClassifierScope(klass, useSiteSession)
+            FirGeneratedClassNestedClassifierScope.create(useSiteSession, klass)
        } else {
            val baseScope = FirNestedClassifierScopeImpl(klass, useSiteSession)
-            if (extensions.any { it.needToGenerateNestedClassifiersInClass(klass) }) {
+            val generatedScope = FirGeneratedClassNestedClassifierScope.create(useSiteSession, klass)
+            if (generatedScope != null) {
                FirCompositeNestedClassifierScope(
-                    listOf(baseScope, FirGeneratedClassNestedClassifierScope(klass, useSiteSession)),
+                    listOf(baseScope, generatedScope),
                    klass,
                    useSiteSession
                )
            } else {
                baseScope
            }
-        }.takeUnless { it.isEmpty() }
+        }?.takeUnless { it.isEmpty() }
    }
 }

@@ -10,25 +10,44 @@ import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.FirLazyValue
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.caches.getValue
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.utils.classId
-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.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.declarations.validate
+import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.utils.addToStdlib.flatGroupBy
 import org.jetbrains.kotlin.utils.addToStdlib.runIf

-class FirGeneratedClassDeclaredMemberScope(
+class FirGeneratedClassDeclaredMemberScope private constructor(
    val useSiteSession: FirSession,
    val firClass: FirClass,
-    needNestedClassifierScope: Boolean
+    needNestedClassifierScope: Boolean,
+    val extensionsByCallableName: Map<Name, List<FirDeclarationGenerationExtension>>,
+    val allCallableNames: Set<Name>
 ) : FirClassDeclaredMemberScope(firClass.classId) {
-    private val extensions: List<FirDeclarationGenerationExtension> = firClass.findGeneratedExtensions(useSiteSession) {
-        needToGenerateAdditionalMembersInClass(it)
+    companion object {
+        fun create(session: FirSession, firClass: FirClass, needNestedClassifierScope: Boolean): FirGeneratedClassDeclaredMemberScope? {
+            val extensionsByCallableName = session.groupExtensionsByName(
+                firClass,
+                nameExtractor = { getCallableNamesForClass(it) },
+                nameTransformer = { it }
+            )
+            val allCallableNames = extensionsByCallableName.keys
+            if (allCallableNames.isEmpty()) return null
+            return FirGeneratedClassDeclaredMemberScope(
+                session,
+                firClass,
+                needNestedClassifierScope,
+                extensionsByCallableName,
+                allCallableNames
+            )
+        }
    }

    private val nestedClassifierScope: FirNestedClassifierScope? = runIf(needNestedClassifierScope) {
@@ -51,38 +70,32 @@ class FirGeneratedClassDeclaredMemberScope(
        generateConstructors()
    }

-    private val callableNamesCache: FirLazyValue<Set<Name>, Nothing?> = firCachesFactory.createLazyValue {
-        extensions.flatMapTo(mutableSetOf()) { it.getCallableNamesForClass(firClass.symbol) }
-    }
-
    // ------------------------------------------ generators ------------------------------------------

    private fun generateMemberFunctions(name: Name): List<FirNamedFunctionSymbol> {
-        return extensions
+        if (name == SpecialNames.INIT) return emptyList()
+        return extensionsByCallableName[name].orEmpty()
            .flatMap { it.generateFunctions(CallableId(firClass.classId, name), firClass.symbol) }
            .onEach { it.fir.validate() }
    }

    private fun generateMemberProperties(name: Name): List<FirPropertySymbol> {
-        return extensions
+        if (name == SpecialNames.INIT) return emptyList()
+        return extensionsByCallableName[name].orEmpty()
            .flatMap { it.generateProperties(CallableId(firClass.classId, name), firClass.symbol) }
            .onEach { it.fir.validate() }
    }

    private fun generateConstructors(): List<FirConstructorSymbol> {
-        val classId = firClass.symbol.classId
-        val callableId = if (classId.isNestedClass) {
-            CallableId(classId.parentClassId!!, classId.shortClassName)
-        } else {
-            CallableId(classId.asSingleFqName().parent(), classId.shortClassName)
-        }
-        return extensions.flatMap { it.generateConstructors(callableId) }.onEach { it.fir.validate() }
+        return extensionsByCallableName[SpecialNames.INIT].orEmpty()
+            .flatMap { it.generateConstructors(firClass.symbol) }
+            .onEach { it.fir.validate() }
    }

    // ------------------------------------------ scope methods ------------------------------------------

    override fun getCallableNames(): Set<Name> {
-        return callableNamesCache.getValue()
+        return allCallableNames
    }

    override fun getClassifierNames(): Set<Name> {
@@ -114,22 +127,51 @@ class FirGeneratedClassDeclaredMemberScope(
    }
 }

-class FirGeneratedClassNestedClassifierScope(
+internal inline fun <T, V> FirSession.groupExtensionsByName(
    klass: FirClass,
-    useSiteSession: FirSession
+    nameExtractor: FirDeclarationGenerationExtension.(FirClassSymbol<*>) -> Set<T>,
+    nameTransformer: (T) -> V
+): Map<V, List<FirDeclarationGenerationExtension>> {
+    val extensions = getExtensionsForClass(klass)
+    val symbol = klass.symbol
+    return extensions.flatGroupBy(
+        keySelector = { extension -> extension.nameExtractor(symbol) },
+        keyTransformer = nameTransformer,
+        valueTransformer = { it }
+    )
+}
+
+internal fun FirSession.getExtensionsForClass(klass: FirClass): List<FirDeclarationGenerationExtension> {
+    val extensions = extensionService.declarationGenerators
+    return if (klass.origin.generated) {
+        val pluginKey = (klass.origin as FirDeclarationOrigin.Plugin).key
+        extensions.filter { it.key == pluginKey }
+    } else {
+        extensions
+    }
+}
+
+class FirGeneratedClassNestedClassifierScope private constructor(
+    useSiteSession: FirSession,
+    klass: FirClass,
+    private val nestedClassifierNames: Set<Name>
 ) : FirNestedClassifierScope(klass, useSiteSession) {
-    private val extensions = klass.findGeneratedExtensions(useSiteSession) { needToGenerateNestedClassifiersInClass(it) }
+    companion object {
+        @OptIn(FirExtensionApiInternals::class)
+        fun create(useSiteSession: FirSession, klass: FirClass): FirGeneratedClassNestedClassifierScope? {
+            val extensions = useSiteSession.getExtensionsForClass(klass)
+            val symbol = klass.symbol
+            val classifierNames = extensions.flatMapTo(mutableSetOf()) { it.nestedClassifierNamesCache.getValue(symbol) }
+            if (classifierNames.isEmpty()) return null
+            return FirGeneratedClassNestedClassifierScope(useSiteSession, klass, classifierNames)
+        }
+    }

    private val nestedClassifierCache: FirCache<Name, FirRegularClassSymbol?, Nothing?> =
        useSiteSession.firCachesFactory.createCache { name, _ ->
            generateNestedClassifier(name)
        }

-    private val nestedClassifiersNames: FirLazyValue<Set<Name>, Nothing?> =
-        useSiteSession.firCachesFactory.createLazyValue {
-            extensions.flatMapTo(mutableSetOf()) { it.getNestedClassifiersNames(klass.symbol) }
-        }
-
    private fun generateNestedClassifier(name: Name): FirRegularClassSymbol? {
        if (klass is FirRegularClass) {
            val companion = klass.companionObjectSymbol
@@ -138,8 +180,9 @@ class FirGeneratedClassNestedClassifierScope(
            }
        }

-        if (name !in getClassifierNames()) return null
-        val generatedClass = useSiteSession.symbolProvider.getClassLikeSymbolByClassId(klass.classId.createNestedClassId(name))
+        if (name !in nestedClassifierNames) return null
+        val generatedClass = useSiteSession.generatedDeclarationsSymbolProvider
+            ?.getClassLikeSymbolByClassId(klass.classId.createNestedClassId(name))
        require(generatedClass is FirRegularClassSymbol?) { "Only regular class are allowed as nested classes" }
        return generatedClass
    }
@@ -149,26 +192,10 @@ class FirGeneratedClassNestedClassifierScope(
    }

    override fun isEmpty(): Boolean {
-        return getClassifierNames().isEmpty()
+        return nestedClassifierNames.isEmpty()
    }

    override fun getClassifierNames(): Set<Name> {
-        return nestedClassifiersNames.getValue()
-    }
-}
-
-
-private inline fun FirClass.findGeneratedExtensions(
-    useSiteSession: FirSession,
-    predicate: FirDeclarationGenerationExtension.(FirClass) -> Boolean
-): List<FirDeclarationGenerationExtension> {
-    val origin = origin
-    val declarationGenerators = useSiteSession.extensionService.declarationGenerators
-    return if (origin is FirDeclarationOrigin.Plugin) {
-        declarationGenerators.filter { it.key == origin.key }.also {
-            require(it.isNotEmpty()) { "Extension for ${origin.key} not found" }
-        }
-    } else {
-        declarationGenerators.filter { it.predicate(this) }
+        return nestedClassifierNames
    }
 }
@@ -13,7 +13,11 @@ import org.jetbrains.kotlin.fir.visitors.FirTransformer
@RequiresOptIn(message = "Should be used just only in resolve processor")
 annotation class AdapterForResolveProcessor

-sealed class FirResolveProcessor(val session: FirSession, val scopeSession: ScopeSession)
+sealed class FirResolveProcessor(val session: FirSession, val scopeSession: ScopeSession) {
+    open fun beforePhase() {}
+
+    open fun afterPhase() {}
+}

 abstract class FirGlobalResolveProcessor(session: FirSession, scopeSession: ScopeSession) : FirResolveProcessor(session, scopeSession) {
    abstract fun process(files: Collection<FirFile>)
@@ -22,6 +22,7 @@ class FirTotalResolveProcessor(session: FirSession, enablePluginPhases: Boolean

    fun process(files: List<FirFile>) {
        for (processor in processors) {
+            processor.beforePhase()
            when (processor) {
                is FirTransformerBasedResolveProcessor -> {
                    for (file in files) {
@@ -32,6 +33,7 @@ class FirTotalResolveProcessor(session: FirSession, enablePluginPhases: Boolean
                    processor.process(files)
                }
            }
+            processor.afterPhase()
        }
    }
 }
@@ -22,11 +22,25 @@ import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.fqName
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.name.FqName

-class FirPluginAnnotationsResolveProcessor(session: FirSession, scopeSession: ScopeSession) : FirTransformerBasedResolveProcessor(session, scopeSession) {
+class FirPluginAnnotationsResolveProcessor(
+    session: FirSession,
+    scopeSession: ScopeSession
+) : FirTransformerBasedResolveProcessor(session, scopeSession) {
    override val transformer = FirPluginAnnotationsResolveTransformer(session, scopeSession)
+
+    @OptIn(FirSymbolProviderInternals::class)
+    override fun beforePhase() {
+        session.generatedDeclarationsSymbolProvider?.disable()
+    }
+
+    @OptIn(FirSymbolProviderInternals::class)
+    override fun afterPhase() {
+        session.generatedDeclarationsSymbolProvider?.enable()
+    }
 }

 class FirPluginAnnotationsResolveTransformer(
@@ -45,3 +45,6 @@ abstract class FirDeclarationPredicateRegistrar {
    abstract fun register(vararg predicates: DeclarationPredicate)
    abstract fun register(predicates: Collection<DeclarationPredicate>)
 }
+
+@RequiresOptIn
+annotation class FirExtensionApiInternals
@@ -12,9 +12,7 @@ import org.jetbrains.kotlin.fir.expressions.FirDelegatedConstructorCall
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.symbols.ensureResolved
-import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.*

 sealed class FirFunctionSymbol<D : FirFunction>(
    override val callableId: CallableId
@@ -57,6 +55,8 @@ class FirIntersectionOverrideFunctionSymbol(
 class FirConstructorSymbol(
    callableId: CallableId
 ) : FirFunctionSymbol<FirConstructor>(callableId) {
+    constructor(classId: ClassId) : this(classId.callableIdForConstructor())
+
    val isPrimary: Boolean
        get() = fir.isPrimary