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

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

@@ -7,11 +7,8 @@ package org.jetbrains.kotlin.fir.plugin.generators
 
 import org.jetbrains.kotlin.GeneratedDeclarationKey
 import org.jetbrains.kotlin.fir.FirSession
-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.*
 import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.plugin.createConstructor
 import org.jetbrains.kotlin.fir.plugin.createMemberFunction
 import org.jetbrains.kotlin.fir.plugin.createNestedClass
@@ -52,7 +49,11 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         return listOf(function.symbol)
     }
 
-    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+    override fun generateNestedClassLikeDeclaration(
+        owner: FirClassSymbol<*>,
+        name: Name,
+        context: NestedClassGenerationContext
+    ): FirClassLikeSymbol<*>? {
         if (matchedClasses.none { it == owner }) return null
         return createNestedClass(owner, name, Key).symbol
     }
@@ -62,7 +63,7 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         return listOf(createConstructor.symbol)
     }
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         return when {
             classSymbol in matchedClasses -> setOf(MATERIALIZE_NAME)
             classSymbol.classId in nestedClassIds -> setOf(SpecialNames.INIT)
@@ -70,7 +71,7 @@ class AdditionalMembersGenerator(session: FirSession) : FirDeclarationGeneration
         }
     }
 
-    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> {
         return if (classSymbol in matchedClasses) setOf(NESTED_NAME) else emptySet()
     }
 

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

@@ -9,11 +9,8 @@ import org.jetbrains.kotlin.GeneratedDeclarationKey
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
-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.*
 import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.plugin.createCompanionObject
 import org.jetbrains.kotlin.fir.plugin.createDefaultPrivateConstructor
 import org.jetbrains.kotlin.fir.plugin.createMemberFunction
@@ -32,7 +29,7 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
         private val FOO_NAME = Name.identifier("foo")
     }
 
-    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name, context: NestedClassGenerationContext): FirClassLikeSymbol<*>? {
         if (name != SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT) return null
         return createCompanionObject(owner, Key).symbol
     }
@@ -51,7 +48,7 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
         return listOf(constructor.symbol)
     }
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         if (classSymbol.classKind != ClassKind.OBJECT) return emptySet()
         if (classSymbol !is FirRegularClassSymbol) return emptySet()
         val classId = classSymbol.classId
@@ -64,7 +61,7 @@ class CompanionGenerator(session: FirSession) : FirDeclarationGenerationExtensio
         }
     }
 
-    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> {
         return if (session.predicateBasedProvider.matches(PREDICATE, classSymbol)) {
             setOf(SpecialNames.DEFAULT_NAME_FOR_COMPANION_OBJECT)
         } else {

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

@@ -8,11 +8,8 @@ package org.jetbrains.kotlin.fir.plugin.generators
 import org.jetbrains.kotlin.GeneratedDeclarationKey
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-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.*
 import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
-import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.plugin.*
 import org.jetbrains.kotlin.fir.resolve.providers.getRegularClassSymbolByClassId
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
@@ -60,7 +57,11 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
         return createTopLevelClass(classId, Key).symbol
     }
 
-    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+    override fun generateNestedClassLikeDeclaration(
+        owner: FirClassSymbol<*>,
+        name: Name,
+        context: NestedClassGenerationContext
+    ): FirClassLikeSymbol<*>? {
         return when (val origin = owner.origin) {
             is FirDeclarationOrigin.Plugin -> when (origin.key) {
                 Key -> generateNestedClass(owner.classId.createNestedClassId(name), owner)
@@ -95,7 +96,7 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
         return listOf(function.symbol)
     }
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         return when (classSymbol.classId) {
             in classIdsForMatchedClasses -> setOf(MATERIALIZE_NAME, SpecialNames.INIT)
             GENERATED_CLASS_ID -> setOf(SpecialNames.INIT)
@@ -103,7 +104,7 @@ class ExternalClassGenerator(session: FirSession) : FirDeclarationGenerationExte
         }
     }
 
-    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> {
         return if (classSymbol.classId == GENERATED_CLASS_ID) {
             return classIdsForMatchedClasses.keys.mapTo(mutableSetOf()) { it.shortClassName }
         } else {

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

@@ -62,7 +62,7 @@ class MembersOfSerializerGenerator(session: FirSession) : FirDeclarationGenerati
         return listOf(function.symbol)
     }
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         return when (classSymbol) {
             in matchedCoreSerializerClasses -> serializeMethodNames.keys
             else -> emptySet()

plugins/kotlinx-serialization/kotlinx-serialization.k2/src/org/jetbrains/kotlinx/serialization/compiler/fir/SerializationFirResolveExtension.kt

@@ -51,7 +51,7 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
         hasFactory && hasAnnotatedFactory
     }
 
-    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> {
         val result = mutableSetOf<Name>()
         with(session) {
             if (classSymbol.shouldHaveGeneratedMethodsInCompanion && !classSymbol.isSerializableObject)
@@ -63,7 +63,11 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
         return result
     }
 
-    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+    override fun generateNestedClassLikeDeclaration(
+        owner: FirClassSymbol<*>,
+        name: Name,
+        context: NestedClassGenerationContext
+    ): FirClassLikeSymbol<*>? {
         if (owner !is FirRegularClassSymbol) return null
         if (!session.predicateBasedProvider.matches(FirSerializationPredicates.annotatedWithSerializableOrMeta, owner)) return null
         return when (name) {
@@ -73,7 +77,7 @@ class SerializationFirResolveExtension(session: FirSession) : FirDeclarationGene
         }
     }
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         val classId = classSymbol.classId
         val result = mutableSetOf<Name>()
 

plugins/lombok/lombok.k2/src/org/jetbrains/kotlin/lombok/k2/generators/BuilderGenerator.kt

@@ -23,10 +23,10 @@ import org.jetbrains.kotlin.fir.declarations.utils.classId
 import org.jetbrains.kotlin.fir.declarations.utils.effectiveVisibility
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
+import org.jetbrains.kotlin.fir.extensions.NestedClassGenerationContext
 import org.jetbrains.kotlin.fir.java.JavaScopeProvider
 import org.jetbrains.kotlin.fir.java.declarations.*
 import org.jetbrains.kotlin.fir.resolve.defaultType
-import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.toEffectiveVisibility
@@ -45,7 +45,6 @@ import org.jetbrains.kotlin.lombok.k2.java.*
 import org.jetbrains.kotlin.lombok.utils.LombokNames
 import org.jetbrains.kotlin.lombok.utils.capitalize
 import org.jetbrains.kotlin.name.CallableId
-import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 
 class BuilderGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
@@ -62,12 +61,12 @@ class BuilderGenerator(session: FirSession) : FirDeclarationGenerationExtension(
     private val functionsCache: FirCache<FirClassSymbol<*>, Map<Name, List<FirJavaMethod>>?, Nothing?> =
         session.firCachesFactory.createCache(::createFunctions)
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         if (!classSymbol.isSuitableJavaClass()) return emptySet()
         return functionsCache.getValue(classSymbol)?.keys.orEmpty()
     }
 
-    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getNestedClassifiersNames(classSymbol: FirClassSymbol<*>, context: NestedClassGenerationContext): Set<Name> {
         if (!classSymbol.isSuitableJavaClass()) return emptySet()
         val name = builderClassCache.getValue(classSymbol)?.name ?: return emptySet()
         return setOf(name)
@@ -78,7 +77,11 @@ class BuilderGenerator(session: FirSession) : FirDeclarationGenerationExtension(
         return functionsCache.getValue(classSymbol)?.get(callableId.callableName).orEmpty().map { it.symbol }
     }
 
-    override fun generateNestedClassLikeDeclaration(owner: FirClassSymbol<*>, name: Name): FirClassLikeSymbol<*>? {
+    override fun generateNestedClassLikeDeclaration(
+        owner: FirClassSymbol<*>,
+        name: Name,
+        context: NestedClassGenerationContext
+    ): FirClassLikeSymbol<*>? {
         if (!owner.isSuitableJavaClass()) return null
         return builderClassCache.getValue(owner)?.symbol
     }

plugins/lombok/lombok.k2/src/org/jetbrains/kotlin/lombok/k2/generators/GetterGenerator.kt

@@ -40,7 +40,7 @@ class GetterGenerator(session: FirSession) : FirDeclarationGenerationExtension(s
     private val cache: FirCache<FirClassSymbol<*>, Map<Name, FirJavaMethod>?, Nothing?> =
         session.firCachesFactory.createCache(::createGetters)
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         if (!classSymbol.isSuitableJavaClass()) return emptySet()
         return cache.getValue(classSymbol)?.keys ?: emptySet()
     }

plugins/lombok/lombok.k2/src/org/jetbrains/kotlin/lombok/k2/generators/LombokConstructorsGenerator.kt

@@ -30,7 +30,7 @@ class LombokConstructorsGenerator(session: FirSession) : FirDeclarationGeneratio
     private val cache: FirCache<FirClassSymbol<*>, Collection<FirFunctionSymbol<*>>?, Nothing?> =
         session.firCachesFactory.createCache(::createConstructors)
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         if (!classSymbol.isSuitableJavaClass()) return emptySet()
         return cache.getValue(classSymbol)?.mapTo(mutableSetOf()) {
             when (it) {

plugins/lombok/lombok.k2/src/org/jetbrains/kotlin/lombok/k2/generators/SetterGenerator.kt

@@ -43,7 +43,7 @@ class SetterGenerator(session: FirSession) : FirDeclarationGenerationExtension(s
     private val cache: FirCache<FirClassSymbol<*>, Map<Name, FirJavaMethod>?, Nothing?> =
         session.firCachesFactory.createCache(::createSetters)
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         if (!classSymbol.isSuitableForSetters()) return emptySet()
         return cache.getValue(classSymbol)?.keys ?: emptySet()
     }

plugins/lombok/lombok.k2/src/org/jetbrains/kotlin/lombok/k2/generators/WithGenerator.kt

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.fir.caches.createCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
 import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.impl.FirResolvedDeclarationStatusImpl
-import org.jetbrains.kotlin.fir.declarations.utils.isFinal
 import org.jetbrains.kotlin.fir.extensions.FirDeclarationGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.MemberGenerationContext
 import org.jetbrains.kotlin.fir.java.declarations.FirJavaField
@@ -25,7 +24,6 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.toEffectiveVisibility
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.isBoolean
 import org.jetbrains.kotlin.lombok.config.AccessLevel
 import org.jetbrains.kotlin.lombok.k2.config.ConeLombokAnnotations.With
 import org.jetbrains.kotlin.lombok.k2.config.LombokService
@@ -34,7 +32,6 @@ import org.jetbrains.kotlin.lombok.utils.collectWithNotNull
 import org.jetbrains.kotlin.lombok.utils.toPropertyNameCapitalized
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.types.typeUtil.isBoolean
 
 class WithGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
     private val lombokService: LombokService
@@ -43,7 +40,7 @@ class WithGenerator(session: FirSession) : FirDeclarationGenerationExtension(ses
     private val cache: FirCache<FirClassSymbol<*>, Map<Name, FirJavaMethod>?, Nothing?> =
         session.firCachesFactory.createCache(::createWith)
 
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         if (!classSymbol.isSuitableJavaClass()) return emptySet()
         return cache.getValue(classSymbol)?.keys ?: emptySet()
     }

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

@@ -105,7 +105,7 @@ class FirParcelizeDeclarationGenerator(session: FirSession) : FirDeclarationGene
     }
 
     @OptIn(SymbolInternals::class)
-    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>): Set<Name> {
+    override fun getCallableNamesForClass(classSymbol: FirClassSymbol<*>, context: MemberGenerationContext): Set<Name> {
         return when {
             classSymbol.fir.modality == Modality.ABSTRACT -> emptySet()
             classSymbol in matchedClasses && classSymbol.fir.modality != Modality.SEALED -> parcelizeMethodsNames