[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 generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> = emptyList()
-    open fun generateProperties(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirPropertySymbol> = emptyList()
+    open fun generateProperties(callableId: CallableId, context: MemberGenerationContext?): List<FirPropertySymbol> = emptyList()
-    open fun generateConstructors(owner: FirClassSymbol<*>): List<FirConstructorSymbol> = 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() }
     }
 

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/AdditionalMembersGenerator.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.annotated
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
@@ -50,7 +51,7 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         matchedClasses.map { it.classId.createNestedClassId(NESTED_NAME) }
     }
 
-    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+    override fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> {
         if (callableId.callableName != MATERIALIZE_NAME) return emptyList()
         val classId = callableId.classId ?: return emptyList()
         val matchedClassSymbol = matchedClasses.firstOrNull { it.classId == classId } ?: return emptyList()
@@ -74,9 +75,10 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         }.symbol
     }
 
-    override fun generateConstructors(owner: FirClassSymbol<*>): List<FirConstructorSymbol> {
+    override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
-        assert(owner.classId in nestedClassIds)
+        val ownerClassId = context.owner.classId
-        return listOf(buildConstructor(owner.classId, isInner = false, Key).symbol)
+        assert(ownerClassId in nestedClassIds)
+        return listOf(buildConstructor(ownerClassId, isInner = false, Key).symbol)
     }
 
     override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/CompanionGenerator.kt

@@ -18,6 +18,7 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
 import org.jetbrains.kotlin.fir.extensions.predicate.annotated
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.moduleData
@@ -69,7 +70,8 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
         return regularClass.symbol
     }
 
-    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+    override fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> {
+        val owner = context?.owner
         if (owner == null || owner.origin != Key.origin) return emptyList()
         if (callableId.callableName != FOO_NAME) return emptyList()
         val function = buildSimpleFunction {
@@ -89,8 +91,8 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
         return listOf(function.symbol)
     }
 
-    override fun generateConstructors(owner: FirClassSymbol<*>): List<FirConstructorSymbol> {
+    override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
-        val constructor = buildConstructor(owner.classId, isInner = false, Key)
+        val constructor = buildConstructor(context.owner.classId, isInner = false, Key)
         return listOf(constructor.symbol)
     }
 

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/ExternalClassGenerator.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.annotated
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
@@ -80,8 +81,8 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
         return buildClass(classId).symbol
     }
 
-    override fun generateConstructors(owner: FirClassSymbol<*>): List<FirConstructorSymbol> {
+    override fun generateConstructors(context: MemberGenerationContext): List<FirConstructorSymbol> {
-        val classId = owner.classId
+        val classId = context.owner.classId
         if (classId != GENERATED_CLASS_ID && classId !in classIdsForMatchedClasses) return emptyList()
         return listOf(buildConstructor(classId, isInner = false, Key).symbol)
     }
@@ -96,8 +97,9 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
     }
 
     @OptIn(SymbolInternals::class)
-    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+    override fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> {
         if (callableId.classId !in classIdsForMatchedClasses || callableId.callableName != MATERIALIZE_NAME) return emptyList()
+        val owner = context?.owner
         require(owner is FirRegularClassSymbol)
         val matchedClassId = owner.fir.matchedClass ?: return emptyList()
         val matchedClassSymbol = session.symbolProvider.getClassLikeSymbolByClassId(matchedClassId) ?: return emptyList()

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/MembersOfSerializerGenerator.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusIm
 import org.jetbrains.kotlin.fir.expressions.builder.buildBlock
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
 import org.jetbrains.kotlin.fir.extensions.predicate.annotated
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.moduleData
@@ -57,7 +58,7 @@ class MembersOfSerializerGenerator(session: FirSession) : FirDeclarationGenerati
         serializableClassIds.associateBy { Name.identifier("serialize${it.shortClassName.identifier}") }
     }
 
-    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+    override fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> {
         val argumentClassId = serializeMethodNames[callableId.callableName] ?: return emptyList()
         val dispatchReceiverClassId = callableId.classId ?: return emptyList()
         val function = buildSimpleFunction {

plugins/fir-plugin-prototype/src/org/jetbrains/kotlin/fir/plugin/generators/TopLevelDeclarationsGenerator.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.builder.buildValueParameter
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.annotated
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
@@ -41,8 +42,8 @@ class TopLevelDeclarationsGenerator(session: FirSession) : FirDeclarationGenerat
         predicateBasedProvider.getSymbolsByPredicate(PREDICATE).filterIsInstance<FirRegularClassSymbol>()
     }
 
-    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+    override fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> {
-        if (owner != null) return emptyList()
+        if (context != null) return emptyList()
         val matchedClassSymbol = findMatchedClassForFunction(callableId) ?: return emptyList()
         val function = buildSimpleFunction {
             resolvePhase = FirResolvePhase.BODY_RESOLVE

plugins/parcelize/parcelize-compiler/parcelize.k2/src/org/jetbrains/kotlin/parcelize/fir/FirParcelizeDeclarationGenerator.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusIm
 import org.jetbrains.kotlin.fir.declarations.utils.modality
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
+import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
 import org.jetbrains.kotlin.fir.extensions.predicate.annotated
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.moduleData
@@ -53,8 +54,8 @@ class FirParcelizeDeclarationGenerator(session: FirSession) : FirDeclarationGene
             .filterIsInstance<FirRegularClassSymbol>()
     }
 
-    override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
+    override fun generateFunctions(callableId: CallableId, context: MemberGenerationContext?): List<FirNamedFunctionSymbol> {
-        if (owner == null) return emptyList()
+        val owner = context?.owner ?: return emptyList()
         require(owner is FirRegularClassSymbol)
         val functionSymbol = when (callableId.callableName) {
             DESCRIBE_CONTENTS_NAME -> {