[FIR] Add prototype of declaring plugin annotations

2020-05-07 12:09:26 +03:00
parent d1b709e632
commit 8f293acb2e
5 changed files with 69 additions and 13 deletions
@@ -6,11 +6,17 @@
 package org.jetbrains.kotlin.fir.extensions

 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name

+typealias AnnotationFqn = FqName
+
 abstract class FirExtensionPoint(val session: FirSession) {
    abstract val name: FirExtensionPointName

+    abstract val annotations: Set<AnnotationFqn>
+    abstract val metaAnnotations: Set<AnnotationFqn>
+
    fun interface Factory<P : FirExtensionPoint> {
        fun create(session: FirSession): P
    }
@@ -5,8 +5,14 @@

 package org.jetbrains.kotlin.fir.extensions

+import com.google.common.collect.LinkedHashMultimap
+import com.google.common.collect.Multimap
+import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.getSymbolByTypeRef
+import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
 import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner
 import org.jetbrains.kotlin.fir.utils.ComponentTypeRegistry
 import kotlin.properties.ReadOnlyProperty
@@ -14,27 +20,62 @@ import kotlin.reflect.KClass
 import kotlin.reflect.KProperty

@Suppress("EXPERIMENTAL_FEATURE_WARNING")
-inline class FirRegisteredExtension<P : FirExtensionPoint>(val extensions: List<P>)
+inline class FirRegisteredExtension<P : FirExtensionPoint>(val extensions: Multimap<AnnotationFqn, P>)

 class FirExtensionPointService(
-    private val session: FirSession
+    val session: FirSession
 ) : ComponentArrayOwner<FirExtensionPoint, FirRegisteredExtension<*>>(), FirSessionComponent {
    companion object : ComponentTypeRegistry<FirExtensionPoint, FirRegisteredExtension<*>>() {
-        inline fun <reified K : FirExtensionPoint, V : FirRegisteredExtension<K>> registeredExtensions(): ReadOnlyProperty<FirExtensionPointService, List<K>> {
-            val accessor = generateAccessor<V, K>(K::class)
-            return object : ReadOnlyProperty<FirExtensionPointService, List<K>> {
-                override fun getValue(thisRef: FirExtensionPointService, property: KProperty<*>): List<K> {
-                    return accessor.getValue(thisRef, property).extensions
+        inline fun <reified P : FirExtensionPoint, V : FirRegisteredExtension<P>> registeredExtensions(): ReadOnlyProperty<FirExtensionPointService, ExtensionsAccessor<P>> {
+            val accessor = generateAccessor<V, P>(P::class)
+            return object : ReadOnlyProperty<FirExtensionPointService, ExtensionsAccessor<P>> {
+                override fun getValue(thisRef: FirExtensionPointService, property: KProperty<*>): ExtensionsAccessor<P> {
+                    return ExtensionsAccessor(thisRef.session, accessor.getValue(thisRef, property).extensions)
                }
            }
        }
+
+        private fun <K, V> createMultimap(): Multimap<K, V> = LinkedHashMultimap.create()
    }

    override val typeRegistry: ComponentTypeRegistry<FirExtensionPoint, FirRegisteredExtension<*>>
        get() = Companion

-    fun <P : FirExtensionPoint> registerExtensions(extensionClass: KClass<P>, extensions: List<FirExtensionPoint.Factory<P>>) {
-        registerComponent(extensionClass, FirRegisteredExtension(extensions.map { it.create(session) }))
+    fun <P : FirExtensionPoint> registerExtensions(extensionClass: KClass<P>, extensionFactories: List<FirExtensionPoint.Factory<P>>) {
+        val extensions = extensionFactories.map { it.create(session) }
+        val map = createMultimap<AnnotationFqn, P>()
+        for (extension in extensions) {
+            _metaAnnotations += extension.metaAnnotations
+            for (annotation in extension.annotations) {
+                _annotations += annotation
+                map.put(annotation, extension)
+            }
+        }
+        registerComponent(extensionClass, FirRegisteredExtension(map))
+    }
+
+    val annotations: Set<AnnotationFqn>
+        get() = _annotations
+    private val _annotations: MutableSet<AnnotationFqn> = mutableSetOf()
+
+    val metaAnnotations: Set<AnnotationFqn>
+        get() = _metaAnnotations
+    private val _metaAnnotations: MutableSet<AnnotationFqn> = mutableSetOf()
+
+    class ExtensionsAccessor<P : FirExtensionPoint>(
+        private val session: FirSession,
+        private val extensions: Multimap<AnnotationFqn, P>
+    ) {
+        fun forDeclaration(declaration: FirAnnotationContainer): Collection<P> {
+            if (declaration.annotations.isEmpty()) return emptySet()
+            val result = mutableSetOf<P>()
+            for (annotation in declaration.annotations) {
+                val symbol = session.firSymbolProvider.getSymbolByTypeRef<FirRegularClassSymbol>(annotation.annotationTypeRef) ?: continue
+                val fqName = symbol.classId.asSingleFqName()
+                result += extensions[fqName]
+            }
+            return result
+        }
    }
 }

@@ -27,7 +27,7 @@ abstract class FirStatusTransformerExtension(session: FirSession) : FirExtension
    fun interface Factory : FirExtensionPoint.Factory<FirStatusTransformerExtension>
 }

-val FirExtensionPointService.statusTransformerExtensions: List<FirStatusTransformerExtension> by FirExtensionPointService.registeredExtensions()
+val FirExtensionPointService.statusTransformerExtensions: FirExtensionPointService.ExtensionsAccessor<FirStatusTransformerExtension> by FirExtensionPointService.registeredExtensions()

 inline fun FirDeclarationStatus.transform(
    visibility: Visibility = this.visibility,
@@ -7,13 +7,19 @@ package org.jetbrains.kotlin.fir.resolve

 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.FirSymbolOwner
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
 import org.jetbrains.kotlin.fir.scopes.impl.nestedClassifierScope
+import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTag
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierLookupTagWithFixedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
+import org.jetbrains.kotlin.fir.types.ConeLookupTagBasedType
+import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -77,3 +83,7 @@ fun FirSymbolProvider.getClassDeclaredCallableSymbols(classId: ClassId, name: Na
    return result
 }

+inline fun <reified T : AbstractFirBasedSymbol<*>> FirSymbolProvider.getSymbolByTypeRef(typeRef: FirTypeRef): T? {
+    val lookupTag = typeRef.coneTypeSafe<ConeLookupTagBasedType>()?.lookupTag ?: return null
+    return getSymbolByLookupTag(lookupTag) as? T
+}
@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.extensions.extensionPointService
 import org.jetbrains.kotlin.fir.extensions.statusTransformerExtensions
-import org.jetbrains.kotlin.fir.firEffectiveVisibility
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.compose
@@ -173,9 +172,9 @@ private class FirStatusResolveTransformer(
        return block.compose()
    }

-    private fun prepareStatus(declaration: FirDeclaration, status: FirDeclarationStatus): FirDeclarationStatus {
+    private fun <D> prepareStatus(declaration: D, status: FirDeclarationStatus): FirDeclarationStatus where D : FirDeclaration, D : FirAnnotationContainer {
        var result = status
-        session.extensionPointService.statusTransformerExtensions.forEach {
+        session.extensionPointService.statusTransformerExtensions.forDeclaration(declaration).forEach {
            result = it.transformStatus(declaration, status)
        }
        return result