[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
2021-10-04 15:44:59 +03:00
parent 20e35167c8
commit 6a14d91dbd
16 changed files with 125 additions and 32 deletions
@@ -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

@@ -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
@@ -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
        }
@@ -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
@@ -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>
@@ -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
@@ -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)
            }
        }
@@ -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,
@@ -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 { predicateBasedProvider.matches(it.predicate, this) }
+        declarationGenerators.filter { it.predicate(this) }
    }
 }
@@ -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>)
+}