[FIR] Pass declared member scope to methods of declaration generation extension

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

@@ -9,6 +9,8 @@ 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.FirTypeScope
+import org.jetbrains.kotlin.fir.scopes.impl.FirClassDeclaredMemberScope
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
@@ -38,9 +40,9 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
     open fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
 
     // 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(owner: FirClassSymbol<*>): List<FirConstructorSymbol> = emptyList()
+    open fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> = emptyList()
+    open fun generateProperties(callableId: CallableId, context: MemberGenerationContext?): List<FirPropertySymbol> = emptyList()
+    open fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> = emptyList()
 
     // Can be called on IMPORTS stage
     open fun hasPackage(packageFqName: FqName): Boolean = false
@@ -79,4 +81,20 @@ abstract class FirDeclarationGenerationExtension(session: FirSession) : FirExten
 
 }
 
+class MemberGenerationContext(
+    val owner: FirClassSymbol<*>,
+    val declaredMemberScope: FirClassDeclaredMemberScope?,
+) {
+    override fun equals(other: Any?): Boolean {
+        if (other !is MemberGenerationContext) {
+            return false
+        }
+        return owner == other.owner
+    }
+
+    override fun hashCode(): Int {
+        return owner.hashCode()
+    }
+}
+
 val FirExtensionService.declarationGenerators: List<FirDeclarationGenerationExtension> by FirExtensionService.registeredExtensions()

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

@@ -98,13 +98,13 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
 
     private fun generateTopLevelFunctions(callableId: CallableId): List<FirNamedFunctionSymbol> {
         return extensionsByTopLevelCallableId.getValue()[callableId].orEmpty()
-            .flatMap { it.generateFunctions(callableId, owner = null) }
+            .flatMap { it.generateFunctions(callableId, context = null) }
             .onEach { it.fir.validate() }
     }
 
     private fun generateTopLevelProperties(callableId: CallableId): List<FirPropertySymbol> {
         return extensionsByTopLevelCallableId.getValue()[callableId].orEmpty()
-            .flatMap { it.generateProperties(callableId, owner = null) }
+            .flatMap { it.generateProperties(callableId, context = null) }
             .onEach { it.fir.validate() }
     }
 

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

@@ -12,6 +12,7 @@ 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.MemberGenerationContext
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 import org.jetbrains.kotlin.fir.scopes.FirNameAwareCompositeScope
@@ -53,7 +54,11 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
     ): FirContainingNamesAwareScope {
         return when {
             klass.origin.generated -> {
-                FirGeneratedClassDeclaredMemberScope.create(useSiteSession, klass, needNestedClassifierScope = true) ?: FirTypeScope.Empty
+                FirGeneratedClassDeclaredMemberScope.create(
+                    useSiteSession,
+                    MemberGenerationContext(klass.symbol, declaredMemberScope = null),
+                    needNestedClassifierScope = true
+                ) ?: FirTypeScope.Empty
             }
             else -> {
                 val baseScope = FirClassDeclaredMemberScopeImpl(
@@ -63,7 +68,11 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
                     existingNames,
                     symbolProvider
                 )
-                val generatedScope = FirGeneratedClassDeclaredMemberScope.create(useSiteSession, klass, needNestedClassifierScope = false)
+                val generatedScope = FirGeneratedClassDeclaredMemberScope.create(
+                    useSiteSession,
+                    MemberGenerationContext(klass.symbol, baseScope),
+                    needNestedClassifierScope = false
+                )
                 if (generatedScope != null) {
                     FirNameAwareCompositeScope(listOf(baseScope, generatedScope))
                 } else {

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

@@ -26,15 +26,19 @@ import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 class FirGeneratedClassDeclaredMemberScope private constructor(
     val useSiteSession: FirSession,
-    val firClass: FirClass,
+    private val generationContext: MemberGenerationContext,
     needNestedClassifierScope: Boolean,
     val extensionsByCallableName: Map<Name, List<FirDeclarationGenerationExtension>>,
     val allCallableNames: Set<Name>
-) : FirClassDeclaredMemberScope(firClass.classId) {
+) : FirClassDeclaredMemberScope(generationContext.owner.classId) {
     companion object {
-        fun create(session: FirSession, firClass: FirClass, needNestedClassifierScope: Boolean): FirGeneratedClassDeclaredMemberScope? {
+        fun create(
+            session: FirSession,
+            generationContext: MemberGenerationContext,
+            needNestedClassifierScope: Boolean
+        ): FirGeneratedClassDeclaredMemberScope? {
             val extensionsByCallableName = session.groupExtensionsByName(
-                firClass,
+                generationContext.owner.fir,
                 nameExtractor = { getCallableNamesForClass(it) },
                 nameTransformer = { it }
             )
@@ -42,7 +46,7 @@ class FirGeneratedClassDeclaredMemberScope private constructor(
             if (allCallableNames.isEmpty()) return null
             return FirGeneratedClassDeclaredMemberScope(
                 session,
-                firClass,
+                generationContext,
                 needNestedClassifierScope,
                 extensionsByCallableName,
                 allCallableNames
@@ -50,6 +54,9 @@ class FirGeneratedClassDeclaredMemberScope private constructor(
         }
     }
 
+    private val firClass: FirClass
+        get() = generationContext.owner.fir
+
     private val nestedClassifierScope: FirNestedClassifierScope? = runIf(needNestedClassifierScope) {
         useSiteSession.nestedClassifierScope(firClass)
     }
@@ -75,20 +82,20 @@ class FirGeneratedClassDeclaredMemberScope private constructor(
     private fun generateMemberFunctions(name: Name): List<FirNamedFunctionSymbol> {
         if (name == SpecialNames.INIT) return emptyList()
         return extensionsByCallableName[name].orEmpty()
-            .flatMap { it.generateFunctions(CallableId(firClass.classId, name), firClass.symbol) }
+            .flatMap { it.generateFunctions(CallableId(firClass.classId, name), generationContext) }
             .onEach { it.fir.validate() }
     }
 
     private fun generateMemberProperties(name: Name): List<FirPropertySymbol> {
         if (name == SpecialNames.INIT) return emptyList()
         return extensionsByCallableName[name].orEmpty()
-            .flatMap { it.generateProperties(CallableId(firClass.classId, name), firClass.symbol) }
+            .flatMap { it.generateProperties(CallableId(firClass.classId, name), generationContext) }
             .onEach { it.fir.validate() }
     }
 
     private fun generateConstructors(): List<FirConstructorSymbol> {
         return extensionsByCallableName[SpecialNames.INIT].orEmpty()
-            .flatMap { it.generateConstructors(firClass.symbol) }
+            .flatMap { it.generateConstructors(generationContext) }
             .onEach { it.fir.validate() }
     }