[FIR] Rework predicates system

Now predicates are split into LookupPredicate and DeclarationPredicate hierarchies. First one allows to perform global search for declarations and second one allows to check if some declaration matches the predicate. Predicates with meta annotations are excluded from LookupPredicates, because it's impossible to create index of annotations with meta-annotations, because they can be located inside binary dependencies (so to achieve this we need to scan the whole classpath). Also only one predicate with meta-annotations is left in DeclarationPredicate hierarchy (AnnotatedWithMeta) ^KT-53874 Fixed ^KT-53590 Fixed
2022-11-23 15:16:44 +02:00
parent 7a9a71a089
commit 246dc985a6
41 changed files with 934 additions and 605 deletions
@@ -95,7 +95,7 @@ internal class LLFirNonUnderContentRootSessionFactory(private val project: Proje
                )
            )

-            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider())
+            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)
            register(FirDependenciesSymbolProvider::class, dependencyProvider)
            register(FirDependenciesSymbolProvider::class, dependencyProvider)
            register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getContainingFile
 import org.jetbrains.kotlin.analysis.providers.KotlinAnnotationsResolver
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.createCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
@@ -23,15 +22,20 @@ import org.jetbrains.kotlin.fir.declarations.FirCallableDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.extensions.AnnotationFqn
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
-import org.jetbrains.kotlin.fir.extensions.predicate.*
+import org.jetbrains.kotlin.fir.extensions.predicate.AbstractPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.LookupPredicate
+import org.jetbrains.kotlin.fir.extensions.predicate.PredicateVisitor
 import org.jetbrains.kotlin.fir.extensions.registeredPluginAnnotations
 import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.fir.resolve.getCorrespondingClassSymbolOrNull
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.classId
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
+import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.name.ClassId
@@ -52,8 +56,8 @@ internal class LLFirIdePredicateBasedProvider(
    private val declarationOwnersCache: FirCache<FirFile, FirOwnersStorage, Nothing?> =
        session.firCachesFactory.createCache { firFile -> FirOwnersStorage.collectOwners(firFile) }

-    override fun getSymbolsByPredicate(predicate: DeclarationPredicate): List<FirBasedSymbol<*>> {
-        val annotations = registeredPluginAnnotations.getAnnotationsForPredicate(predicate)
+    override fun getSymbolsByPredicate(predicate: LookupPredicate): List<FirBasedSymbol<*>> {
+        val annotations = predicate.annotations
        val annotatedDeclarations = annotations
            .asSequence()
            .flatMap { annotationsResolver.declarationsByAnnotation(ClassId.topLevel(it)) }
@@ -99,80 +103,83 @@ internal class LLFirIdePredicateBasedProvider(
        }
    }

-    override fun matches(predicate: DeclarationPredicate, declaration: FirDeclaration): Boolean {
-        return predicate.accept(matcher, declaration)
+    override fun matches(predicate: AbstractPredicate<*>, declaration: FirDeclaration): Boolean {
+        return when (predicate) {
+            is DeclarationPredicate -> predicate.accept(declarationPredicateMatcher, declaration)
+            is LookupPredicate -> predicate.accept(lookupPredicateMatcher, declaration)
+        }
    }

-    private val matcher: Matcher = Matcher()
+    private val declarationPredicateMatcher = Matcher<DeclarationPredicate>()
+    private val lookupPredicateMatcher = Matcher<LookupPredicate>()

-    private inner class Matcher : DeclarationPredicateVisitor<Boolean, FirDeclaration>() {
-        override fun visitPredicate(predicate: DeclarationPredicate, data: FirDeclaration): Boolean {
-            error("When overrides for all possible DeclarationPredicate subtypes are implemented, " +
-                          "this method should never be called, but it was called with $predicate")
+    private inner class Matcher<P : AbstractPredicate<P>> : PredicateVisitor<P, Boolean, FirDeclaration>() {
+        override fun visitPredicate(predicate: AbstractPredicate<P>, data: FirDeclaration): Boolean {
+            error(
+                "When overrides for all possible DeclarationPredicate subtypes are implemented, " +
+                        "this method should never be called, but it was called with $predicate"
+            )
        }

-        override fun visitAnd(predicate: DeclarationPredicate.And, data: FirDeclaration): Boolean {
+        override fun visitAnd(predicate: AbstractPredicate.And<P>, data: FirDeclaration): Boolean {
            return predicate.a.accept(this, data) && predicate.b.accept(this, data)
        }

-        override fun visitOr(predicate: DeclarationPredicate.Or, data: FirDeclaration): Boolean {
+        override fun visitOr(predicate: AbstractPredicate.Or<P>, data: FirDeclaration): Boolean {
            return predicate.a.accept(this, data) || predicate.b.accept(this, data)
        }

-        override fun visitAnnotatedWith(predicate: AnnotatedWith, data: FirDeclaration): Boolean {
+        override fun visitAnnotatedWith(predicate: AbstractPredicate.AnnotatedWith<P>, data: FirDeclaration): Boolean {
            return annotationsOnDeclaration(data).any { it in predicate.annotations }
        }

-        override fun visitAncestorAnnotatedWith(predicate: AncestorAnnotatedWith, data: FirDeclaration): Boolean {
+        override fun visitAncestorAnnotatedWith(predicate: AbstractPredicate.AncestorAnnotatedWith<P>, data: FirDeclaration): Boolean {
            return annotationsOnOuterDeclarations(data).any { it in predicate.annotations }
        }

-        override fun visitMetaAnnotatedWith(predicate: MetaAnnotatedWith, data: FirDeclaration): Boolean {
-            return metaAnnotationsOnDeclaration(data).any { it in predicate.metaAnnotations }
+        override fun visitMetaAnnotatedWith(predicate: AbstractPredicate.MetaAnnotatedWith<P>, data: FirDeclaration): Boolean {
+            val visited = mutableSetOf<FirRegularClassSymbol>()
+            return data.annotations.any { annotation ->
+                annotation.markedWithMetaAnnotation(predicate.metaAnnotations, visited)
+            }
        }

-        override fun visitAncestorMetaAnnotatedWith(predicate: AncestorMetaAnnotatedWith, data: FirDeclaration): Boolean {
-            return metaAnnotationsOnOuterDeclarations(data).any { it in predicate.metaAnnotations }
+        private fun FirAnnotation.markedWithMetaAnnotation(
+            metaAnnotations: Set<AnnotationFqn>,
+            visited: MutableSet<FirRegularClassSymbol>
+        ): Boolean {
+            val symbol = this.getCorrespondingClassSymbolOrNull(session) ?: return false
+            if (!visited.add(symbol)) return false
+            if (symbol.classId.asSingleFqName() in metaAnnotations) return true
+            if (symbol.resolvedAnnotationsWithClassIds.any { it.markedWithMetaAnnotation(metaAnnotations, visited) }) return true
+            return false
        }

-        override fun visitParentAnnotatedWith(predicate: ParentAnnotatedWith, data: FirDeclaration): Boolean {
+        override fun visitParentAnnotatedWith(predicate: AbstractPredicate.ParentAnnotatedWith<P>, data: FirDeclaration): Boolean {
            val parent = data.directParentDeclaration ?: return false
-            val parentPredicate = AnnotatedWith(predicate.annotations)
+            val parentPredicate = DeclarationPredicate.AnnotatedWith(predicate.annotations)

-            return parentPredicate.accept(this, parent)
+            return parentPredicate.accept(declarationPredicateMatcher, parent)
        }

-        override fun visitHasAnnotatedWith(predicate: HasAnnotatedWith, data: FirDeclaration): Boolean {
-            val childPredicate = AnnotatedWith(predicate.annotations)
-
-            return data.anyDirectChildDeclarationMatches(childPredicate)
-        }
-
-        override fun visitParentMetaAnnotatedWith(predicate: ParentMetaAnnotatedWith, data: FirDeclaration): Boolean {
-            val parent = data.directParentDeclaration ?: return false
-            val parentPredicate = MetaAnnotatedWith(predicate.annotations)
-
-            return parentPredicate.accept(this, parent)
-        }
-
-        override fun visitHasMetaAnnotatedWith(predicate: HasMetaAnnotatedWith, data: FirDeclaration): Boolean {
-            val childPredicate = MetaAnnotatedWith(predicate.annotations)
+        override fun visitHasAnnotatedWith(predicate: AbstractPredicate.HasAnnotatedWith<P>, data: FirDeclaration): Boolean {
+            val childPredicate = DeclarationPredicate.AnnotatedWith(predicate.annotations)

            return data.anyDirectChildDeclarationMatches(childPredicate)
        }

        private val FirDeclaration.directParentDeclaration: FirDeclaration?
            get() = getOwnersOfDeclaration(this)?.lastOrNull()?.fir
+    }

-        private fun FirDeclaration.anyDirectChildDeclarationMatches(childPredicate: DeclarationPredicate): Boolean {
-            var result = false
+    private fun FirDeclaration.anyDirectChildDeclarationMatches(childPredicate: DeclarationPredicate): Boolean {
+        var result = false

-            this.forEachDirectChildDeclaration {
-                result = result || childPredicate.accept(this@Matcher, it)
-            }
-
-            return result
+        this.forEachDirectChildDeclaration {
+            result = result || childPredicate.accept(declarationPredicateMatcher, it)
        }
+
+        return result
    }

    private fun annotationsOnDeclaration(declaration: FirDeclaration): Set<AnnotationFqn> {
@@ -193,24 +200,9 @@ internal class LLFirIdePredicateBasedProvider(
        return psiAnnotations.map { it.asSingleFqName() }.toSet()
    }

-    private fun metaAnnotationsOnDeclaration(declaration: FirDeclaration): Set<AnnotationFqn> {
-        val directAnnotations = annotationsOnDeclaration(declaration)
-        val metaAnnotations = directAnnotations
-            .asSequence()
-            .mapNotNull { declarationProvider.getAllClassesByClassId(ClassId.topLevel(it)).singleOrNull() }
-            .flatMap { annotationsResolver.annotationsOnDeclaration(it) }
-            .toSet()
-
-        return metaAnnotations.map { it.asSingleFqName() }.toSet()
-    }
-
    private fun annotationsOnOuterDeclarations(declaration: FirDeclaration): Set<AnnotationFqn> {
        return getOwnersOfDeclaration(declaration)?.flatMap { annotationsOnDeclaration(it.fir) }.orEmpty().toSet()
    }
-
-    private fun metaAnnotationsOnOuterDeclarations(declaration: FirDeclaration): Set<AnnotationFqn> {
-        return getOwnersOfDeclaration(declaration)?.flatMap { metaAnnotationsOnDeclaration(it.fir) }.orEmpty().toSet()
-    }
 }

 private class FirOwnersStorage(private val declarationToOwner: Map<FirDeclaration, List<FirBasedSymbol<*>>>) {
@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
 import org.jetbrains.kotlin.analysis.providers.KotlinAnnotationsResolver
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.caches.*
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.extensions.AbstractFirRegisteredPluginAnnotations
 import org.jetbrains.kotlin.fir.extensions.AnnotationFqn
 import org.jetbrains.kotlin.name.ClassId
@@ -26,25 +25,14 @@ internal class LLFirIdeRegisteredPluginAnnotations(
        get() = allAnnotationsCache.getValue()

    private val allAnnotationsCache: FirLazyValue<Set<AnnotationFqn>, Nothing?> = session.firCachesFactory.createLazyValue {
-        // at this point, both metaAnnotations and annotationsFromPlugins should be collected
-        val result = metaAnnotations.flatMapTo(mutableSetOf()) { getAnnotationsWithMetaAnnotation(it) }
-
-        if (result.isEmpty()) {
-            annotationsFromPlugins
-        } else {
-            result += annotationsFromPlugins
-            result
-        }
+        // at this point, annotationsFromPlugins should be collected
+        annotationsFromPlugins
    }

    // MetaAnnotation -> Annotations
    private val annotationsWithMetaAnnotationCache: FirCache<AnnotationFqn, Set<AnnotationFqn>, Nothing?> =
        session.firCachesFactory.createCache { metaAnnotation -> collectAnnotationsWithMetaAnnotation(metaAnnotation) }

-    override fun getAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Collection<AnnotationFqn> {
-        return annotationsWithMetaAnnotationCache.getValue(metaAnnotation)
-    }
-
    private fun collectAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Set<FqName> {
        val annotatedDeclarations = annotationsResolver.declarationsByAnnotation(ClassId.topLevel(metaAnnotation))

@@ -59,8 +47,4 @@ internal class LLFirIdeRegisteredPluginAnnotations(
    override fun saveAnnotationsFromPlugin(annotations: Collection<AnnotationFqn>) {
        annotationsFromPlugins += annotations
    }
-
-    override fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection<FirRegularClass>) {
-        error("This method should never be called in IDE mode")
-    }
 }
@@ -214,7 +214,7 @@ internal object LLFirSessionFactory {
            // We need FirRegisteredPluginAnnotations during extensions' registration process
            val annotationsResolver = project.createAnnotationResolver(contentScope)
            register(FirRegisteredPluginAnnotations::class, LLFirIdeRegisteredPluginAnnotations(this@session, annotationsResolver))
-            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider())
+            register(FirPredicateBasedProvider::class, FirEmptyPredicateBasedProvider)

            val dependencyProvider = LLFirDependentModuleProvidersByProviders(this) {
                // <all libraries scope> - <current library scope>
@@ -24,7 +24,7 @@ internal class LLFirDesignatedAnnotationsResolveTransformed(
        if (!designationIterator.hasNext()) {
            val declaration = designation.declaration
            if (declaration is FirRegularClass || declaration is FirTypeAlias) {
-                declaration.transform<FirDeclaration, Mode>(this, Mode.RegularAnnotations)
+                declaration.transform<FirDeclaration, Nothing?>(this, null)
            }
            return
        }
@@ -63,5 +63,4 @@ internal class LLFirDesignatedAnnotationsResolveTransformed(
        // todo add proper check that COMPILER_REQUIRED_ANNOTATIONS are resolved
 //        checkNestedDeclarationsAreResolved(declaration)
    }
-
 }