From 8f293acb2e1d88599aac4fce080e98f4e6525828 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 7 May 2020 12:09:26 +0300
Subject: [PATCH] [FIR] Add prototype of declaring plugin annotations

---
 .../fir/extensions/FirExtensionPoint.kt       |  6 ++
 .../extensions/FirExtensionPointService.kt    | 59 ++++++++++++++++---
 .../FirStatusTransfomerExtension.kt           |  2 +-
 .../kotlin/fir/resolve/FirSymbolProvider.kt   | 10 ++++
 .../FirStatusResolveTransformer.kt            |  5 +-
 5 files changed, 69 insertions(+), 13 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt
index ebcd9c64e22..f5dec2ecebe 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPoint.kt
@@ -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
     }
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt
index 66ae2bf202f..3e65f8fb63a 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionPointService.kt
@@ -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
+        }
     }
 }
 
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt
index 231c384a460..93f81429d1e 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirStatusTransfomerExtension.kt
@@ -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,
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSymbolProvider.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSymbolProvider.kt
index 77677e846a6..8453d2d3961 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSymbolProvider.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSymbolProvider.kt
@@ -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
+}
\ No newline at end of file
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt
index 2e8d3549143..a13b1e23d71 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt
@@ -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