[FIR] Pass context to all relevant methods of FirDeclarationGenerationExtension

^KT-53470 Fixed

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

@@ -9,7 +9,9 @@ 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.scopes.FirContainingNamesAwareScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassDeclaredMemberScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirNestedClassifierScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
@@ -38,7 +40,11 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
      */
     open fun generateTopLevelClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
 
-    open fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? = null
+    open fun generateNestedClassLikeDeclaration(
+        owner: FirClassSymbol<*>,
+        name: Name,
+        context: NestedClassGenerationContext
+    ): FirClassLikeSymbol<*>? = null
 
     // Can be called on STATUS stage
     open fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> = emptyList()
@@ -57,8 +63,8 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
      * 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 getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> = emptySet()
+    open fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> = emptySet()
     open fun getTopLevelCallableIds(): Set<CallableId> = emptySet()
     open fun getTopLevelClassIds(): Set<ClassId> = emptySet()
 
@@ -67,9 +73,9 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
     // ----------------------------------- internal utils -----------------------------------
 
     @FirExtensionApiInternals
-    val nestedClassifierNamesCache: FirCache<FirClassSymbol<*>, Set<Name>, Nothing?> =
-        session.firCachesFactory.createCache { symbol, _ ->
-            getNestedClassifiersNames(symbol)
+    val nestedClassifierNamesCache: FirCache<FirClassSymbol<*>, Set<Name>, NestedClassGenerationContext> =
+        session.firCachesFactory.createCache { symbol, context ->
+            getNestedClassifiersNames(symbol, context)
         }
 
     @FirExtensionApiInternals
@@ -82,19 +88,42 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
 
 }
 
-class MemberGenerationContext(
+typealias MemberGenerationContext = DeclarationGenerationContext.Member
+typealias NestedClassGenerationContext = DeclarationGenerationContext.Nested
+
+sealed class DeclarationGenerationContext<T : FirContainingNamesAwareScope>(
     val owner: FirClassSymbol<*>,
-    val declaredMemberScope: FirClassDeclaredMemberScope?,
+    val declaredScope: T?,
 ) {
+    // is needed for `hashCode` implementation
+    protected abstract val kind: Int
+
+    class Member(
+        owner: FirClassSymbol<*>,
+        declaredScope: FirClassDeclaredMemberScope?,
+    ) : DeclarationGenerationContext<FirClassDeclaredMemberScope>(owner, declaredScope) {
+        override val kind: Int
+            get() = 1
+    }
+
+    class Nested(
+        owner: FirClassSymbol<*>,
+        declaredScope: FirNestedClassifierScope?,
+    ) : DeclarationGenerationContext<FirNestedClassifierScope>(owner, declaredScope) {
+        override val kind: Int
+            get() = 2
+    }
+
     override fun equals(other: Any?): Boolean {
-        if (other !is MemberGenerationContext) {
+        if (this.javaClass !== other?.javaClass) {
             return false
         }
+        require(other is DeclarationGenerationContext<*>)
         return owner == other.owner
     }
 
     override fun hashCode(): Int {
-        return owner.hashCode()
+        return owner.hashCode() + kind
     }
 }
 

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

@@ -58,7 +58,7 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
             origin.generated -> {
                 FirGeneratedClassDeclaredMemberScope.create(
                     useSiteSession,
-                    MemberGenerationContext(klass.symbol, declaredMemberScope = null),
+                    MemberGenerationContext(klass.symbol, declaredScope = null),
                     needNestedClassifierScope = true
                 ) ?: FirTypeScope.Empty
             }
@@ -92,10 +92,10 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
 
     private fun createNestedClassifierScope(klass: FirClass): FirNestedClassifierScope? {
         return if (klass.origin.generated) {
-            FirGeneratedClassNestedClassifierScope.create(useSiteSession, klass)
+            FirGeneratedClassNestedClassifierScope.create(useSiteSession, klass, baseScope = null)
         } else {
             val baseScope = FirNestedClassifierScopeImpl(klass, useSiteSession)
-            val generatedScope = FirGeneratedClassNestedClassifierScope.create(useSiteSession, klass)
+            val generatedScope = FirGeneratedClassNestedClassifierScope.create(useSiteSession, klass, baseScope)
             if (generatedScope != null) {
                 FirCompositeNestedClassifierScope(
                     listOf(baseScope, generatedScope),

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

@@ -40,7 +40,7 @@ class FirGeneratedClassDeclaredMemberScope private constructor(
         ): FirGeneratedClassDeclaredMemberScope? {
             val extensionsByCallableName = session.groupExtensionsByName(
                 generationContext.owner.fir,
-                nameExtractor = { getCallableNamesForClass(it) },
+                nameExtractor = { getCallableNamesForClass(it, generationContext) },
                 nameTransformer = { it }
             )
             val allCallableNames = extensionsByCallableName.keys
@@ -162,17 +162,22 @@ class FirGeneratedClassNestedClassifierScope private constructor(
     useSiteSession: FirSession,
     klass: FirClass,
     private val extensionsByName: Map<Name, List<FirDeclarationGenerationExtension>>,
+    private val context: NestedClassGenerationContext
 ) : FirNestedClassifierScope(klass, useSiteSession) {
     companion object {
         @OptIn(FirExtensionApiInternals::class)
-        fun create(useSiteSession: FirSession, klass: FirClass): FirGeneratedClassNestedClassifierScope? {
+        fun create(
+            useSiteSession: FirSession,
+            klass: FirClass,
+            baseScope: FirNestedClassifierScope?
+        ): FirGeneratedClassNestedClassifierScope? {
             val symbol = klass.symbol
-
+            val context = NestedClassGenerationContext(klass.symbol, baseScope)
             val extensionsByName = useSiteSession.getExtensionsForClass(klass).flatGroupBy {
-                it.nestedClassifierNamesCache.getValue(symbol)
+                it.nestedClassifierNamesCache.getValue(symbol, context)
             }
             if (extensionsByName.isEmpty()) return null
-            return FirGeneratedClassNestedClassifierScope(useSiteSession, klass, extensionsByName)
+            return FirGeneratedClassNestedClassifierScope(useSiteSession, klass, extensionsByName, context)
         }
     }
 
@@ -192,7 +197,7 @@ class FirGeneratedClassNestedClassifierScope private constructor(
         val extensions = extensionsByName[name] ?: return null
 
         val generatedClasses = extensions.mapNotNull { extension ->
-            extension.generateNestedClassLikeDeclaration(klass.symbol, name)?.also { symbol ->
+            extension.generateNestedClassLikeDeclaration(klass.symbol, name, context)?.also { symbol ->
                 symbol.fir.ownerGenerator = extension
             }
         }