[FIR] Change predicate providing API for FIR extensions

Predicate declaration was split into two parts:
- registering any number of predicates which are used in specific extension
- special checks for declaration applicability for each specific extension

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
@@ -27,6 +28,9 @@ 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
     open fun generateClassLikeDeclaration(classId: ClassId): FirClassLikeSymbol<*>? = null
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt

@@ -47,6 +47,9 @@ abstract class FirPredicateBasedProvider : FirSessionComponent {
     abstract fun registerGeneratedDeclaration(declaration: FirAnnotatedDeclaration, owner: FirAnnotatedDeclaration)
 
     abstract fun matches(predicate: DeclarationPredicate, declaration: FirAnnotatedDeclaration): Boolean
+    fun matches(predicates: List<DeclarationPredicate>, declaration: FirAnnotatedDeclaration): Boolean {
+        return predicates.any { matches(it, declaration) }
+    }
 }
 
 @NoMutableState

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirRegisteredPluginAnnotations.kt

@@ -68,9 +68,25 @@ private class FirRegisteredPluginAnnotationsImpl(session: FirSession) : FirRegis
 
     @PluginServicesInitialization
     override fun initialize() {
+        val registrar = object : FirDeclarationPredicateRegistrar() {
+            val predicates = mutableListOf<DeclarationPredicate>()
+            override fun register(vararg predicates: DeclarationPredicate) {
+                this.predicates += predicates
+            }
+
+            override fun register(predicates: Collection<DeclarationPredicate>) {
+                this.predicates += predicates
+            }
+        }
+
         for (extension in session.extensionService.getAllExtensions()) {
             if (extension !is FirPredicateBasedExtension) continue
-            val predicate = extension.predicate
+            with(extension) {
+                registrar.registerPredicates()
+            }
+        }
+
+        for (predicate in registrar.predicates) {
             annotations += predicate.annotations
             metaAnnotations += predicate.metaAnnotations
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransformerExtension.kt

@@ -22,6 +22,8 @@ abstract class FirStatusTransformerExtension(session: FirSession) : FirPredicate
 
     final override val extensionType: KClass<out FirExtension> = FirStatusTransformerExtension::class
 
+    abstract fun needTransformStatus(declaration: FirAnnotatedDeclaration): Boolean
+
     protected open fun transformStatus(
         status: FirDeclarationStatus,
         declaration: FirAnnotatedDeclaration

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirSupertypeGenerationExtension.kt

@@ -21,6 +21,8 @@ abstract class FirSupertypeGenerationExtension(session: FirSession) : FirPredica
 
     final override val extensionType: KClass<out FirExtension> = FirSupertypeGenerationExtension::class
 
+    abstract fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean
+
     abstract fun computeAdditionalSupertypes(
         classLikeDeclaration: FirClassLikeDeclaration,
         resolvedSupertypes: List<FirResolvedTypeRef>

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolver.kt

@@ -53,7 +53,7 @@ class FirStatusResolver(
         if (extensionStatusTransformers.isEmpty()) return status
         val declaration = this
         return extensionStatusTransformers.fold(status) { acc, it ->
-            if (session.predicateBasedProvider.matches(it.predicate, declaration)) {
+            if (it.needTransformStatus(declaration)) {
                 it.operation(acc)
             } else {
                 acc

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSupertypesResolution.kt

@@ -372,9 +372,8 @@ open class FirSupertypeResolverVisitor(
 
     private fun addSupertypesFromExtensions(klass: FirClassLikeDeclaration, supertypeRefs: MutableList<FirResolvedTypeRef>) {
         if (supertypeGenerationExtensions.isEmpty()) return
-        val provider = session.predicateBasedProvider
         for (extension in supertypeGenerationExtensions) {
-            if (provider.matches(extension.predicate, klass)) {
+            if (extension.needTransformSupertypes(klass)) {
                 supertypeRefs += extension.computeAdditionalSupertypes(klass, supertypeRefs)
             }
         }

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

@@ -14,10 +14,8 @@ 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.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.resolve.declaredMemberScopeProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.scopes.FirCompositeScope
 import org.jetbrains.kotlin.fir.scopes.FirContainingNamesAwareScope
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.name.ClassId
@@ -68,7 +66,7 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
                     existingNames,
                     symbolProvider
                 )
-                if (extensions.any { useSiteSession.predicateBasedProvider.matches(it.predicate, klass) }) {
+                if (extensions.any { it.needToGenerateAdditionalMembersInClass(klass) }) {
                     FirCompositeScope(
                         listOf(baseScope, FirGeneratedClassDeclaredMemberScope(useSiteSession, klass, needNestedClassifierScope = false))
                     )
@@ -88,7 +86,7 @@ class FirDeclaredMemberScopeProvider(val useSiteSession: FirSession) : FirSessio
             FirGeneratedClassNestedClassifierScope(klass, useSiteSession)
         } else {
             val baseScope = FirNestedClassifierScopeImpl(klass, useSiteSession)
-            if (extensions.any { useSiteSession.predicateBasedProvider.matches(it.predicate, klass) }) {
+            if (extensions.any { it.needToGenerateNestedClassifiersInClass(klass) }) {
                 FirCompositeNestedClassifierScope(
                     listOf(baseScope, FirGeneratedClassNestedClassifierScope(klass, useSiteSession)),
                     klass,

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

@@ -16,7 +16,6 @@ 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.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.*
@@ -29,7 +28,7 @@ class FirGeneratedClassDeclaredMemberScope(
     val firClass: FirClass,
     needNestedClassifierScope: Boolean
 ) : FirClassDeclaredMemberScope() {
-    private val extensions: List<FirDeclarationGenerationExtension> = firClass.findGeneratedExtensions(useSiteSession)
+    private val extensions: List<FirDeclarationGenerationExtension> = firClass.findGeneratedExtensions(useSiteSession) { needToGenerateAdditionalMembersInClass(it) }
     private val nestedClassifierScope: FirNestedClassifierScope? = runIf(needNestedClassifierScope) {
         useSiteSession.nestedClassifierScope(firClass)
     }
@@ -113,7 +112,7 @@ class FirGeneratedClassNestedClassifierScope(
     klass: FirClass,
     useSiteSession: FirSession
 ) : FirNestedClassifierScope(klass, useSiteSession) {
-    private val extensions = klass.findGeneratedExtensions(useSiteSession)
+    private val extensions = klass.findGeneratedExtensions(useSiteSession) { needToGenerateNestedClassifiersInClass(it) }
 
     private val nestedClassifierCache: FirCache<Name, FirRegularClassSymbol?, Nothing?> =
         useSiteSession.firCachesFactory.createCache { name, _ ->
@@ -146,7 +145,10 @@ class FirGeneratedClassNestedClassifierScope(
 }
 
 
-private fun FirClass.findGeneratedExtensions(useSiteSession: FirSession): List<FirDeclarationGenerationExtension> {
+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) {
@@ -154,7 +156,6 @@ private fun FirClass.findGeneratedExtensions(useSiteSession: FirSession): List<F
             require(it.isNotEmpty()) { "Extension for ${origin.key} not found" }
         }
     } else {
-        val predicateBasedProvider = useSiteSession.predicateBasedProvider
+        declarationGenerators.filter { it.predicate(this) }
-        declarationGenerators.filter { predicateBasedProvider.matches(it.predicate, this) }
     }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirExtension.kt

@@ -34,9 +34,14 @@ abstract class FirExtension(val session: FirSession) {
 }
 
 abstract class FirPredicateBasedExtension(session: FirSession) : FirExtension(session) {
-    abstract val predicate: DeclarationPredicate
+    abstract fun FirDeclarationPredicateRegistrar.registerPredicates()
 }
 
 data class FirExtensionPointName(val name: Name) {
     constructor(name: String) : this(Name.identifier(name))
 }
+
+abstract class FirDeclarationPredicateRegistrar {
+    abstract fun register(vararg predicates: DeclarationPredicate)
+    abstract fun register(predicates: Collection<DeclarationPredicate>)
+}

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

@@ -8,20 +8,22 @@ package org.jetbrains.kotlin.fir.plugin
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
 import org.jetbrains.kotlin.fir.extensions.FirStatusTransformerExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.hasOrUnder
 import org.jetbrains.kotlin.fir.extensions.predicate.metaHasOrUnder
 import org.jetbrains.kotlin.fir.extensions.predicate.or
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.transform
 import org.jetbrains.kotlin.name.FqName
 
 class AllOpenStatusTransformer(session: FirSession) : FirStatusTransformerExtension(session) {
     companion object {
         private val ALL_OPEN = FqName("org.jetbrains.kotlin.fir.plugin.AllOpen")
+        private val PREDICATE: DeclarationPredicate = hasOrUnder(ALL_OPEN) or metaHasOrUnder(ALL_OPEN)
     }
 
     override fun transformStatus(status: FirDeclarationStatus, declaration: FirAnnotatedDeclaration): FirDeclarationStatus {
@@ -29,7 +31,13 @@ class AllOpenStatusTransformer(session: FirSession) : FirStatusTransformerExtens
         return status.transform(modality = Modality.OPEN)
     }
 
-    override val predicate: DeclarationPredicate = hasOrUnder(ALL_OPEN) or metaHasOrUnder(ALL_OPEN)
+    override fun needTransformStatus(declaration: FirAnnotatedDeclaration): Boolean {
+        return session.predicateBasedProvider.matches(PREDICATE, declaration)
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
 
     override val key: FirPluginKey
         get() = AllOpenPluginKey

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

@@ -8,9 +8,11 @@ package org.jetbrains.kotlin.fir.plugin
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
 import org.jetbrains.kotlin.fir.extensions.FirSupertypeGenerationExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.has
+import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.classId
@@ -25,6 +27,8 @@ import org.jetbrains.kotlin.name.Name
 class AllOpenSupertypeGenerator(session: FirSession) : FirSupertypeGenerationExtension(session) {
     companion object {
         private val myInterfaceClassId = ClassId(FqName("foo"), Name.identifier("MyInterface"))
+        private val PREDICATE: DeclarationPredicate = has("D".fqn())
+
     }
 
     override fun computeAdditionalSupertypes(
@@ -42,5 +46,11 @@ class AllOpenSupertypeGenerator(session: FirSession) : FirSupertypeGenerationExt
     override val key: FirPluginKey
         get() = AllOpenPluginKey
 
-    override val predicate: DeclarationPredicate = has("D".fqn())
+    override fun needTransformSupertypes(declaration: FirClassLikeDeclaration): Boolean {
+        return session.predicateBasedProvider.matches(PREDICATE, declaration)
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
 }

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.expressions.arguments
+import org.jetbrains.kotlin.fir.extensions.FirDeclarationPredicateRegistrar
 import org.jetbrains.kotlin.fir.extensions.FirStatusTransformerExtension
 import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.hasOrUnder
@@ -36,6 +37,8 @@ class AllOpenVisibilityTransformer(session: FirSession) : FirStatusTransformerEx
         private val InternalName = Name.identifier("Internal")
         private val PrivateName = Name.identifier("Private")
         private val ProtectedName = Name.identifier("Protected")
+
+        private val PREDICATE: DeclarationPredicate = hasOrUnder(AllPublicClassId.asSingleFqName())
     }
 
     override fun transformStatus(status: FirDeclarationStatus, declaration: FirAnnotatedDeclaration): FirDeclarationStatus {
@@ -68,7 +71,13 @@ class AllOpenVisibilityTransformer(session: FirSession) : FirStatusTransformerEx
         }
     }
 
-    override val predicate: DeclarationPredicate = hasOrUnder(AllPublicClassId.asSingleFqName())
+    override fun needTransformStatus(declaration: FirAnnotatedDeclaration): Boolean {
+        return session.predicateBasedProvider.matches(PREDICATE, declaration)
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
 
     override val key: FirPluginKey
         get() = AllOpenPluginKey

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

@@ -11,11 +11,10 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.buildPrimaryConstructor
 import org.jetbrains.kotlin.fir.declarations.builder.buildRegularClass
-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.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.has
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
@@ -25,8 +24,6 @@ import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.scopes.kotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.*
-import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -46,6 +43,8 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
         private val FOO_PACKAGE = FqName.topLevel(Name.identifier("foo"))
         private val GENERATED_CLASS_ID = ClassId(FOO_PACKAGE, Name.identifier("AllOpenGenerated"))
         private val MATERIALIZE_NAME = Name.identifier("materialize")
+
+        private val PREDICATE: DeclarationPredicate = has("B".fqn())
     }
 
     object Key : FirPluginKey() {
@@ -56,7 +55,7 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
 
     private val predicateBasedProvider = session.predicateBasedProvider
     private val matchedClasses by lazy {
-        predicateBasedProvider.getSymbolsByPredicate(predicate).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
+        predicateBasedProvider.getSymbolsByPredicate(PREDICATE).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
     }
     private val classIdsForMatchedClasses: Map<ClassId, FirRegularClassSymbol> by lazy {
         matchedClasses.associateBy {
@@ -165,8 +164,17 @@ class AllOpenClassGenerator(session: FirSession) : FirDeclarationGenerationExten
     override val key: FirPluginKey
         get() = Key
 
-    override val predicate: DeclarationPredicate
+    override fun needToGenerateAdditionalMembersInClass(klass: FirClass): Boolean {
-        get() = has("B".fqn())
+        return false
+    }
+
+    override fun needToGenerateNestedClassifiersInClass(klass: FirClass): Boolean {
+        return false
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
 }
 
 private object MatchedClassAttributeKey : FirDeclarationDataKey()

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

@@ -10,10 +10,12 @@ import org.jetbrains.kotlin.descriptors.EffectiveVisibility
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 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.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.has
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
@@ -34,11 +36,13 @@ class AllOpenMembersGenerator(session: FirSession) : FirDeclarationGenerationExt
     companion object {
         private val MATERIALIZE_NAME = Name.identifier("materialize")
         private val NESTED_NAME = Name.identifier("Nested")
+
+        private val PREDICATE: DeclarationPredicate = has("C".fqn())
     }
 
     private val predicateBasedProvider = session.predicateBasedProvider
     private val matchedClasses by lazy {
-        predicateBasedProvider.getSymbolsByPredicate(predicate).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
+        predicateBasedProvider.getSymbolsByPredicate(PREDICATE).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
     }
 
     override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
@@ -87,6 +91,16 @@ class AllOpenMembersGenerator(session: FirSession) : FirDeclarationGenerationExt
 
     override val key: FirPluginKey
         get() = Key
-    override val predicate: DeclarationPredicate
+
-        get() = has("C".fqn())
+    override fun needToGenerateAdditionalMembersInClass(klass: FirClass): Boolean {
+        return session.predicateBasedProvider.matches(PREDICATE, klass)
+    }
+
+    override fun needToGenerateNestedClassifiersInClass(klass: FirClass): Boolean {
+        return session.predicateBasedProvider.matches(PREDICATE, klass)
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
 }

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

@@ -9,10 +9,12 @@ import org.jetbrains.kotlin.descriptors.EffectiveVisibility
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.builder.buildSimpleFunction
 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.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.fir.extensions.predicate.has
 import org.jetbrains.kotlin.fir.extensions.predicateBasedProvider
@@ -30,9 +32,13 @@ import org.jetbrains.kotlin.name.Name
  * Generates `dummyClassName(value: ClassName): String` function for each class annotated with @A
  */
 class AllOpenTopLevelDeclarationsGenerator(session: FirSession) : FirDeclarationGenerationExtension(session) {
+    companion object {
+        private val PREDICATE: DeclarationPredicate = has("A".fqn())
+    }
+
     private val predicateBasedProvider = session.predicateBasedProvider
     private val matchedClasses by lazy {
-        predicateBasedProvider.getSymbolsByPredicate(predicate).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
+        predicateBasedProvider.getSymbolsByPredicate(PREDICATE).map { it.symbol }.filterIsInstance<FirRegularClassSymbol>()
     }
 
     override fun generateFunctions(callableId: CallableId, owner: FirClassSymbol<*>?): List<FirNamedFunctionSymbol> {
@@ -87,7 +93,15 @@ class AllOpenTopLevelDeclarationsGenerator(session: FirSession) : FirDeclaration
     override val key: AllOpenPluginKey
         get() = AllOpenPluginKey
 
-    override val predicate: DeclarationPredicate
+    override fun needToGenerateAdditionalMembersInClass(klass: FirClass): Boolean {
-        get() = has("A".fqn())
+        return false
+    }
 
+    override fun needToGenerateNestedClassifiersInClass(klass: FirClass): Boolean {
+        return false
+    }
+
+    override fun FirDeclarationPredicateRegistrar.registerPredicates() {
+        register(PREDICATE)
+    }
 }