[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

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

@@ -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()

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirExtensionDeclarationsSymbolProvider.kt

@@ -0,0 +1,136 @@
+/*
+ * 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.FirSessionComponent
+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,
+    private val extensions: List<FirDeclarationGenerationExtension>
+) : FirSymbolProvider(session), FirSessionComponent {
+    companion object {
+        fun create(session: FirSession): FirExtensionDeclarationsSymbolProvider? {
+            val extensions = session.extensionService.declarationGenerators
+            if (extensions.isEmpty()) return null
+            return FirExtensionDeclarationsSymbolProvider(session, session.firCachesFactory, extensions)
+        }
+    }
+
+    // ------------------------------------------ caches ------------------------------------------
+
+    private val classCache: FirCache<ClassId, FirClassLikeSymbol<*>?, Nothing?> = cachesFactory.createCache { classId, _ ->
+        generateClassLikeDeclaration(classId)
+    }
+
+    private val functionCache: FirCache<CallableId, List<FirNamedFunctionSymbol>, Nothing?> = cachesFactory.createCache { callableId, _ ->
+        generateTopLevelFunctions(callableId)
+    }
+
+    private val propertyCache: FirCache<CallableId, List<FirPropertySymbol>, Nothing?> = cachesFactory.createCache { callableId, _ ->
+        generateTopLevelProperties(callableId)
+    }
+
+    private val packageCache: FirCache<FqName, Boolean, Nothing?> = cachesFactory.createCache { packageFqName, _ ->
+        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 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()
+            else -> error("Multiple plugins generated classes with same classId $classId\n${generatedClasses.joinToString("\n") { it.fir.render() }}")
+        }
+    }
+
+    private fun generateTopLevelFunctions(callableId: CallableId): List<FirNamedFunctionSymbol> {
+        return extensionsByTopLevelCallableId.getValue()[callableId].orEmpty()
+            .flatMap { it.generateFunctions(callableId, owner = null) }
+            .onEach { it.fir.validate() }
+    }
+
+    private fun generateTopLevelProperties(callableId: CallableId): List<FirPropertySymbol> {
+        return extensionsByTopLevelCallableId.getValue()[callableId].orEmpty()
+            .flatMap { it.generateProperties(callableId, owner = null) }
+            .onEach { it.fir.validate() }
+    }
+
+    private fun hasPackage(packageFqName: FqName): Boolean {
+        return extensions.any { it.hasPackage(packageFqName) }
+    }
+
+    // ------------------------------------------ provider methods ------------------------------------------
+
+    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
+        return classCache.getValue(classId)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
+        val callableId = CallableId(packageFqName, name)
+        destination += functionCache.getValue(callableId)
+        destination += propertyCache.getValue(callableId)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
+        destination += functionCache.getValue(CallableId(packageFqName, name))
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
+        destination += propertyCache.getValue(CallableId(packageFqName, name))
+    }
+
+    override fun getPackage(fqName: FqName): FqName? {
+        return fqName.takeIf { packageCache.getValue(fqName, null) }
+    }
+}

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirSwitchableExtensionDeclarationsSymbolProvider.kt

@@ -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()

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirDeclaredMemberScopeProvider.kt

@@ -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() }
     }
 }
 

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirGeneratedClassDeclaredMemberScope.kt

@@ -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
     }
 }