[FIR] Introduce FirPredicateBasedProvider

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.FirSessionBase
 import org.jetbrains.kotlin.fir.FirSessionProvider
 import org.jetbrains.kotlin.fir.analysis.CheckersComponent
 import org.jetbrains.kotlin.fir.extensions.FirOldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
 import org.jetbrains.kotlin.fir.resolve.FirProvider
 import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
@@ -126,6 +127,11 @@ class FirLibrarySession private constructor(
             FirOldExtensionsService(this)
         )
 
+        registerComponent(
+            FirPredicateBasedProvider::class,
+            FirPredicateBasedProvider.create(this)
+        )
+
         sessionProvider.sessionCache[moduleInfo] = this
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirModuleBasedSession.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.fir.extensions.FirOldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.resolve.*
 import org.jetbrains.kotlin.fir.resolve.impl.*
 import org.jetbrains.kotlin.fir.scopes.impl.FirDeclaredMemberScopeProvider
@@ -18,6 +19,7 @@ abstract class FirModuleBasedSession(override val moduleInfo: ModuleInfo, sessio
         registerComponent(FirTypeResolver::class, FirTypeResolverImpl(this))
         registerComponent(FirDeclaredMemberScopeProvider::class, FirDeclaredMemberScopeProvider())
         registerComponent(FirOldExtensionsService::class, FirOldExtensionsService(this))
+        registerComponent(FirPredicateBasedProvider::class, FirPredicateBasedProvider.create(this))
     }
 }
 

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.extensions
+
+import com.google.common.collect.ArrayListMultimap
+import com.google.common.collect.Multimap
+import kotlinx.collections.immutable.PersistentList
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+
+abstract class FirPredicateBasedProvider : FirSessionComponent {
+    companion object {
+        fun create(session: FirSession): FirPredicateBasedProvider {
+            return FirPredicateBasedProviderImpl(session)
+        }
+    }
+
+    abstract fun getSymbolsByPredicate(predicate: DeclarationPredicate): List<FirAnnotatedDeclaration>
+
+    abstract fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>)
+}
+
+private class FirPredicateBasedProviderImpl(private val session: FirSession) : FirPredicateBasedProvider() {
+    private val registeredPluginAnnotations = session.registeredPluginAnnotations
+    private val cache = Cache()
+
+    override fun getSymbolsByPredicate(predicate: DeclarationPredicate): List<FirAnnotatedDeclaration> {
+        val annotations = registeredPluginAnnotations.getAnnotationsForPredicate(predicate)
+        if (annotations.isEmpty()) return emptyList()
+        return annotations.flatMap { cache.declarationByAnnotation[it] + cache.declarationsUnderAnnotated[it] }.filter {
+            predicate.match(it, cache.ownersForDeclaration.getValue(it))
+        }
+    }
+
+    override fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>) {
+        cache.ownersForDeclaration[declaration] = owners
+        registerOwnersDeclarations(declaration, owners)
+
+        if (declaration.annotations.isEmpty()) return
+        val matchingAnnotations = declaration.annotations.mapNotNull { it.fqName(session) }
+            .filter { it in registeredPluginAnnotations.annotations }
+        if (matchingAnnotations.isEmpty()) return
+        matchingAnnotations.forEach { cache.declarationByAnnotation.put(it, declaration) }
+        cache.annotationsOfDeclaration.putAll(declaration, matchingAnnotations)
+    }
+
+    private fun registerOwnersDeclarations(declaration: FirAnnotatedDeclaration, owners: PersistentList<FirAnnotatedDeclaration>) {
+        val lastOwner = owners.lastOrNull() ?: return
+        val annotationsFromLastOwner = cache.annotationsOfDeclaration[lastOwner]
+        val annotationsFromPreviousOwners = cache.parentAnnotationsOfDeclaration[lastOwner]
+
+        val allParentDeclarations = annotationsFromLastOwner + annotationsFromPreviousOwners
+        allParentDeclarations.forEach { cache.declarationsUnderAnnotated.put(it, declaration) }
+        cache.parentAnnotationsOfDeclaration.putAll(declaration, allParentDeclarations)
+    }
+
+    private class Cache {
+        val declarationByAnnotation: Multimap<AnnotationFqn, FirAnnotatedDeclaration> = ArrayListMultimap.create()
+        val annotationsOfDeclaration: Multimap<FirAnnotatedDeclaration, AnnotationFqn> = ArrayListMultimap.create()
+
+        val declarationsUnderAnnotated: Multimap<AnnotationFqn, FirAnnotatedDeclaration> = ArrayListMultimap.create()
+        val parentAnnotationsOfDeclaration: Multimap<FirAnnotatedDeclaration, AnnotationFqn> = ArrayListMultimap.create()
+
+        val ownersForDeclaration: MutableMap<FirAnnotatedDeclaration, PersistentList<FirAnnotatedDeclaration>> = mutableMapOf()
+    }
+}
+
+val FirSession.predicateBasedProvider: FirPredicateBasedProvider by FirSession.sessionComponentAccessor()

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

@@ -0,0 +1,75 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.extensions
+
+import com.google.common.collect.LinkedHashMultimap
+import com.google.common.collect.Multimap
+import org.jetbrains.kotlin.descriptors.ClassKind
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
+
+abstract class FirRegisteredPluginAnnotations : FirSessionComponent {
+    companion object {
+        fun create(): FirRegisteredPluginAnnotations {
+            return FirRegisteredPluginAnnotationsImpl()
+        }
+    }
+
+    abstract val annotations: Set<AnnotationFqn>
+    abstract val metaAnnotations: Set<AnnotationFqn>
+    abstract fun getAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Collection<AnnotationFqn>
+
+    abstract fun registerAnnotations(annotations: Collection<AnnotationFqn>)
+    abstract fun registerMetaAnnotations(metaAnnotations: Collection<AnnotationFqn>)
+    abstract fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection<FirRegularClass>)
+
+    abstract fun getAnnotationsForPredicate(predicate: DeclarationPredicate): Set<AnnotationFqn>
+}
+
+private class FirRegisteredPluginAnnotationsImpl : FirRegisteredPluginAnnotations() {
+    override val annotations: MutableSet<AnnotationFqn> = mutableSetOf()
+    override val metaAnnotations: MutableSet<AnnotationFqn> = mutableSetOf()
+
+    // MetaAnnotation -> Annotations
+    private val userDefinedAnnotations: Multimap<AnnotationFqn, AnnotationFqn> = LinkedHashMultimap.create()
+
+    private val annotationsForPredicateCache: MutableMap<DeclarationPredicate, Set<AnnotationFqn>> = mutableMapOf()
+
+    override fun getAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Collection<AnnotationFqn> {
+        return userDefinedAnnotations[metaAnnotation]
+    }
+
+    override fun registerAnnotations(annotations: Collection<AnnotationFqn>) {
+        this.annotations += annotations
+    }
+
+    override fun registerMetaAnnotations(metaAnnotations: Collection<AnnotationFqn>) {
+        this.metaAnnotations += metaAnnotations
+    }
+
+    override fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection<FirRegularClass>) {
+        require(annotationClasses.all { it.classKind == ClassKind.ANNOTATION_CLASS })
+        val annotations = annotationClasses.map { it.symbol.classId.asSingleFqName() }
+        registerAnnotations(annotations)
+        userDefinedAnnotations.putAll(metaAnnotation, annotations)
+    }
+
+    override fun getAnnotationsForPredicate(predicate: DeclarationPredicate): Set<AnnotationFqn> {
+        return annotationsForPredicateCache.computeIfAbsent(predicate, ::collectAnnotations)
+    }
+
+    private fun collectAnnotations(predicate: DeclarationPredicate): Set<AnnotationFqn> {
+        if (predicate.metaAnnotations.isEmpty()) return predicate.annotations
+        val result = predicate.metaAnnotations.flatMapTo(mutableSetOf()) { getAnnotationsWithMetaAnnotation(it) }
+        if (result.isEmpty()) return predicate.annotations
+        result += predicate.annotations
+        return result
+    }
+}
+
+val FirSession.registeredPluginAnnotations: FirRegisteredPluginAnnotations by FirSession.sessionComponentAccessor()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAbstractAnnotationResolveTransformer.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.transformers.plugin
 
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
@@ -19,7 +20,7 @@ import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.compose
 
-internal abstract class FirAbstractAnnotationResolveTransformer<D>(
+internal abstract class FirAbstractAnnotationResolveTransformer<D, S>(
     protected val session: FirSession,
     protected val scopeSession: ScopeSession
 ) : FirDefaultTransformer<D>() {
@@ -30,57 +31,92 @@ internal abstract class FirAbstractAnnotationResolveTransformer<D>(
     override fun transformFile(file: FirFile, data: D): CompositeTransformResult<FirDeclaration> {
         return withScopeCleanup(towerScope.scopes) {
             towerScope.addScopes(createImportingScopes(file, session, scopeSession))
+            val state = beforeChildren(file)
             file.transformDeclarations(this, data)
-            file.transformAnnotations(this, data).compose()
+            afterChildren(state)
+            transformAnnotatedDeclaration(file, data)
         }
     }
 
     override fun transformProperty(property: FirProperty, data: D): CompositeTransformResult<FirDeclaration> {
-        return property.transformAnnotations(this, data).compose()
+        return transformAnnotatedDeclaration(property, data)
     }
 
     override fun transformRegularClass(
         regularClass: FirRegularClass,
         data: D
     ): CompositeTransformResult<FirStatement> {
-        regularClass.transformDeclarations(this, data)
+        @Suppress("UNCHECKED_CAST")
-        regularClass.transformCompanionObject(this, data)
+        return transformAnnotatedDeclaration(regularClass, data).also {
-        regularClass.transformSuperTypeRefs(this, data)
+            val state = beforeChildren(regularClass)
-        return regularClass.transformAnnotations(this, data).compose()
+            regularClass.transformDeclarations(this, data)
+            regularClass.transformCompanionObject(this, data)
+            regularClass.transformSuperTypeRefs(this, data)
+            afterChildren(state)
+        } as CompositeTransformResult<FirStatement>
     }
 
     override fun transformSimpleFunction(
         simpleFunction: FirSimpleFunction,
         data: D
     ): CompositeTransformResult<FirDeclaration> {
-        simpleFunction.transformValueParameters(this, data)
+        return transformAnnotatedDeclaration(simpleFunction, data).also {
-        return simpleFunction.transformAnnotations(this, data).compose()
+            val state = beforeChildren(simpleFunction)
+            simpleFunction.transformValueParameters(this, data)
+            afterChildren(state)
+        }
     }
 
     override fun transformConstructor(
         constructor: FirConstructor,
         data: D
     ): CompositeTransformResult<FirDeclaration> {
-        constructor.transformValueParameters(this, data)
+        return transformAnnotatedDeclaration(constructor, data).also {
-        return constructor.transformAnnotations(this, data).compose()
+            val state = beforeChildren(constructor)
+            constructor.transformValueParameters(this, data)
+            afterChildren(state)
+        }
     }
 
     override fun transformValueParameter(
         valueParameter: FirValueParameter,
         data: D
     ): CompositeTransformResult<FirStatement> {
-        return valueParameter.transformAnnotations(this, data).compose()
+        @Suppress("UNCHECKED_CAST")
+        return transformAnnotatedDeclaration(valueParameter, data) as CompositeTransformResult<FirStatement>
     }
 
     override fun transformTypeAlias(typeAlias: FirTypeAlias, data: D): CompositeTransformResult<FirDeclaration> {
-        return typeAlias.transformAnnotations(this, data).compose()
+        return transformAnnotatedDeclaration(typeAlias, data)
     }
 
     override fun transformTypeRef(typeRef: FirTypeRef, data: D): CompositeTransformResult<FirTypeRef> {
-        return typeRef.transformAnnotations(this, data).compose()
+        @Suppress("UNCHECKED_CAST")
+        return transformAnnotationContainer(typeRef, data) as CompositeTransformResult<FirTypeRef>
+    }
+
+    override fun transformAnnotatedDeclaration(
+        annotatedDeclaration: FirAnnotatedDeclaration,
+        data: D
+    ): CompositeTransformResult<FirDeclaration> {
+        @Suppress("UNCHECKED_CAST")
+        return transformAnnotationContainer(annotatedDeclaration, data) as CompositeTransformResult<FirDeclaration>
+    }
+
+    override fun transformAnnotationContainer(
+        annotationContainer: FirAnnotationContainer,
+        data: D
+    ): CompositeTransformResult<FirAnnotationContainer> {
+        return annotationContainer.transformAnnotations(this, data).compose()
     }
 
     override fun <E : FirElement> transformElement(element: E, data: D): CompositeTransformResult<E> {
         return element.compose()
     }
+
+    protected open fun beforeChildren(declaration: FirAnnotatedDeclaration): S? {
+        return null
+    }
+
+    protected open fun afterChildren(state: S?) {}
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirPluginAnnotationsResolveTransformer.kt

@@ -7,16 +7,15 @@ package org.jetbrains.kotlin.fir.resolve.transformers.plugin
 
 import com.google.common.collect.LinkedHashMultimap
 import com.google.common.collect.Multimap
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.persistentListOf
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirStatement
-import org.jetbrains.kotlin.fir.extensions.AnnotationFqn
+import org.jetbrains.kotlin.fir.extensions.*
-import org.jetbrains.kotlin.fir.extensions.extensionsService
-import org.jetbrains.kotlin.fir.extensions.fqName
-import org.jetbrains.kotlin.fir.extensions.hasExtensions
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirAbstractPhaseTransformer
 import org.jetbrains.kotlin.fir.resolve.transformers.FirImportResolveTransformer
@@ -42,11 +41,13 @@ class FirPluginAnnotationsResolveTransformer(
     override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirFile> {
         checkSessionConsistency(file)
         if (!extensionPointService.hasExtensions) return file.compose()
+        val registeredPluginAnnotations = file.session.registeredPluginAnnotations
         file.replaceResolvePhase(FirResolvePhase.ANNOTATIONS_FOR_PLUGINS)
         val newAnnotations = file.resolveAnnotations(extensionPointService.annotations, extensionPointService.metaAnnotations)
         if (!newAnnotations.isEmpty) {
             for (metaAnnotation in newAnnotations.keySet()) {
                 extensionPointService.registerUserDefinedAnnotation(metaAnnotation, newAnnotations[metaAnnotation])
+                registeredPluginAnnotations.registerUserDefinedAnnotation(metaAnnotation, newAnnotations[metaAnnotation])
             }
             val newAnnotationsFqns = newAnnotations.values().mapTo(mutableSetOf()) { it.symbol.classId.asSingleFqName() }
             file.resolveAnnotations(newAnnotationsFqns, emptySet())
@@ -80,15 +81,27 @@ private class FirPartialImportResolveTransformer(
 private class FirAnnotationResolveTransformer(
     session: FirSession,
     scopeSession: ScopeSession
-) : FirAbstractAnnotationResolveTransformer<Multimap<AnnotationFqn, FirRegularClass>>(session, scopeSession) {
+) : FirAbstractAnnotationResolveTransformer<Multimap<AnnotationFqn, FirRegularClass>, PersistentList<FirAnnotatedDeclaration>>(session, scopeSession) {
     var metaAnnotations: Set<AnnotationFqn> = emptySet()
-
     private val typeResolverTransformer: FirSpecificTypeResolverTransformer = FirSpecificTypeResolverTransformer(
         towerScope,
         session,
         errorTypeAsResolved = false
     )
 
+    private var owners: PersistentList<FirAnnotatedDeclaration> = persistentListOf()
+
+    override fun beforeChildren(declaration: FirAnnotatedDeclaration): PersistentList<FirAnnotatedDeclaration>? {
+        val current = owners
+        owners = owners.add(declaration)
+        return current
+    }
+
+    override fun afterChildren(state: PersistentList<FirAnnotatedDeclaration>?) {
+        requireNotNull(state)
+        owners = state
+    }
+
     override fun transformAnnotationCall(
         annotationCall: FirAnnotationCall,
         data: Multimap<AnnotationFqn, FirRegularClass>
@@ -109,4 +122,13 @@ private class FirAnnotationResolveTransformer(
             }
         }
     }
+
+    override fun transformAnnotatedDeclaration(
+        annotatedDeclaration: FirAnnotatedDeclaration,
+        data: Multimap<AnnotationFqn, FirRegularClass>
+    ): CompositeTransformResult<FirDeclaration> {
+        return super.transformAnnotatedDeclaration(annotatedDeclaration, data).also {
+            session.predicateBasedProvider.registerAnnotatedDeclaration(annotatedDeclaration, owners)
+        }
+    }
 }