diff --git a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt b/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt index 78039e04868..e17d3e203b3 100644 --- a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt +++ b/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 } diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirModuleBasedSession.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirModuleBasedSession.kt index 77dff5e96f2..ae1441b1d14 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirModuleBasedSession.kt +++ b/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)) } } diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirPredicateBasedProvider.kt new file mode 100644 index 00000000000..b5e2926f190 --- /dev/null +++ b/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 + + abstract fun registerAnnotatedDeclaration(declaration: FirAnnotatedDeclaration, owners: PersistentList) +} + +private class FirPredicateBasedProviderImpl(private val session: FirSession) : FirPredicateBasedProvider() { + private val registeredPluginAnnotations = session.registeredPluginAnnotations + private val cache = Cache() + + override fun getSymbolsByPredicate(predicate: DeclarationPredicate): List { + 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) { + 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) { + 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 = ArrayListMultimap.create() + val annotationsOfDeclaration: Multimap = ArrayListMultimap.create() + + val declarationsUnderAnnotated: Multimap = ArrayListMultimap.create() + val parentAnnotationsOfDeclaration: Multimap = ArrayListMultimap.create() + + val ownersForDeclaration: MutableMap> = mutableMapOf() + } +} + +val FirSession.predicateBasedProvider: FirPredicateBasedProvider by FirSession.sessionComponentAccessor() \ No newline at end of file diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirRegisteredPluginAnnotations.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirRegisteredPluginAnnotations.kt new file mode 100644 index 00000000000..632f31351f3 --- /dev/null +++ b/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 + abstract val metaAnnotations: Set + abstract fun getAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Collection + + abstract fun registerAnnotations(annotations: Collection) + abstract fun registerMetaAnnotations(metaAnnotations: Collection) + abstract fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection) + + abstract fun getAnnotationsForPredicate(predicate: DeclarationPredicate): Set +} + +private class FirRegisteredPluginAnnotationsImpl : FirRegisteredPluginAnnotations() { + override val annotations: MutableSet = mutableSetOf() + override val metaAnnotations: MutableSet = mutableSetOf() + + // MetaAnnotation -> Annotations + private val userDefinedAnnotations: Multimap = LinkedHashMultimap.create() + + private val annotationsForPredicateCache: MutableMap> = mutableMapOf() + + override fun getAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Collection { + return userDefinedAnnotations[metaAnnotation] + } + + override fun registerAnnotations(annotations: Collection) { + this.annotations += annotations + } + + override fun registerMetaAnnotations(metaAnnotations: Collection) { + this.metaAnnotations += metaAnnotations + } + + override fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection) { + 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 { + return annotationsForPredicateCache.computeIfAbsent(predicate, ::collectAnnotations) + } + + private fun collectAnnotations(predicate: DeclarationPredicate): Set { + 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() \ No newline at end of file diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAbstractAnnotationResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAbstractAnnotationResolveTransformer.kt index 2b69a4861df..ce4200b4a86 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAbstractAnnotationResolveTransformer.kt +++ b/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( +internal abstract class FirAbstractAnnotationResolveTransformer( protected val session: FirSession, protected val scopeSession: ScopeSession ) : FirDefaultTransformer() { @@ -30,57 +31,92 @@ internal abstract class FirAbstractAnnotationResolveTransformer( override fun transformFile(file: FirFile, data: D): CompositeTransformResult { 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 { - return property.transformAnnotations(this, data).compose() + return transformAnnotatedDeclaration(property, data) } override fun transformRegularClass( regularClass: FirRegularClass, data: D ): CompositeTransformResult { - regularClass.transformDeclarations(this, data) - regularClass.transformCompanionObject(this, data) - regularClass.transformSuperTypeRefs(this, data) - return regularClass.transformAnnotations(this, data).compose() + @Suppress("UNCHECKED_CAST") + return transformAnnotatedDeclaration(regularClass, data).also { + val state = beforeChildren(regularClass) + regularClass.transformDeclarations(this, data) + regularClass.transformCompanionObject(this, data) + regularClass.transformSuperTypeRefs(this, data) + afterChildren(state) + } as CompositeTransformResult } override fun transformSimpleFunction( simpleFunction: FirSimpleFunction, data: D ): CompositeTransformResult { - simpleFunction.transformValueParameters(this, data) - return simpleFunction.transformAnnotations(this, data).compose() + return transformAnnotatedDeclaration(simpleFunction, data).also { + val state = beforeChildren(simpleFunction) + simpleFunction.transformValueParameters(this, data) + afterChildren(state) + } } override fun transformConstructor( constructor: FirConstructor, data: D ): CompositeTransformResult { - constructor.transformValueParameters(this, data) - return constructor.transformAnnotations(this, data).compose() + return transformAnnotatedDeclaration(constructor, data).also { + val state = beforeChildren(constructor) + constructor.transformValueParameters(this, data) + afterChildren(state) + } } override fun transformValueParameter( valueParameter: FirValueParameter, data: D ): CompositeTransformResult { - return valueParameter.transformAnnotations(this, data).compose() + @Suppress("UNCHECKED_CAST") + return transformAnnotatedDeclaration(valueParameter, data) as CompositeTransformResult } override fun transformTypeAlias(typeAlias: FirTypeAlias, data: D): CompositeTransformResult { - return typeAlias.transformAnnotations(this, data).compose() + return transformAnnotatedDeclaration(typeAlias, data) } override fun transformTypeRef(typeRef: FirTypeRef, data: D): CompositeTransformResult { - return typeRef.transformAnnotations(this, data).compose() + @Suppress("UNCHECKED_CAST") + return transformAnnotationContainer(typeRef, data) as CompositeTransformResult + } + + override fun transformAnnotatedDeclaration( + annotatedDeclaration: FirAnnotatedDeclaration, + data: D + ): CompositeTransformResult { + @Suppress("UNCHECKED_CAST") + return transformAnnotationContainer(annotatedDeclaration, data) as CompositeTransformResult + } + + override fun transformAnnotationContainer( + annotationContainer: FirAnnotationContainer, + data: D + ): CompositeTransformResult { + return annotationContainer.transformAnnotations(this, data).compose() } override fun transformElement(element: E, data: D): CompositeTransformResult { return element.compose() } + + protected open fun beforeChildren(declaration: FirAnnotatedDeclaration): S? { + return null + } + + protected open fun afterChildren(state: S?) {} } \ No newline at end of file diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirPluginAnnotationsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirPluginAnnotationsResolveTransformer.kt index 9bfce171db7..1f2c46c2bd1 100644 --- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirPluginAnnotationsResolveTransformer.kt +++ b/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.extensionsService -import org.jetbrains.kotlin.fir.extensions.fqName -import org.jetbrains.kotlin.fir.extensions.hasExtensions +import org.jetbrains.kotlin.fir.extensions.* 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 { 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>(session, scopeSession) { +) : FirAbstractAnnotationResolveTransformer, PersistentList>(session, scopeSession) { var metaAnnotations: Set = emptySet() - private val typeResolverTransformer: FirSpecificTypeResolverTransformer = FirSpecificTypeResolverTransformer( towerScope, session, errorTypeAsResolved = false ) + private var owners: PersistentList = persistentListOf() + + override fun beforeChildren(declaration: FirAnnotatedDeclaration): PersistentList? { + val current = owners + owners = owners.add(declaration) + return current + } + + override fun afterChildren(state: PersistentList?) { + requireNotNull(state) + owners = state + } + override fun transformAnnotationCall( annotationCall: FirAnnotationCall, data: Multimap @@ -109,4 +122,13 @@ private class FirAnnotationResolveTransformer( } } } + + override fun transformAnnotatedDeclaration( + annotatedDeclaration: FirAnnotatedDeclaration, + data: Multimap + ): CompositeTransformResult { + return super.transformAnnotatedDeclaration(annotatedDeclaration, data).also { + session.predicateBasedProvider.registerAnnotatedDeclaration(annotatedDeclaration, owners) + } + } } \ No newline at end of file