From 78b8a7b38867c645dd0ab1e1813641e295ff836f Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 7 May 2020 14:46:16 +0300
Subject: [PATCH] [FIR] Store extensions by plugin key

---
 .../kotlin/fir/extensions/FirExtensionPoint.kt    |  2 ++
 .../fir/extensions/FirExtensionPointService.kt    | 15 +++++++++++++--
 .../fir/declarations/FirDeclarationOrigin.kt      |  4 ++--
 3 files changed, 17 insertions(+), 4 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 20441288af4..fa6d516e9c7 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,6 +6,7 @@
 package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import kotlin.reflect.KClass
@@ -20,6 +21,7 @@ abstract class FirExtensionPoint(val session: FirSession) {
     abstract val metaAnnotations: Map<AnnotationFqn, MetaAnnotationMode>
 
     abstract val mode: Mode
+    abstract val key: FirPluginKey
 
     internal abstract val extensionType: KClass<out FirExtensionPoint>
 
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 2c58e7e7869..8851de0c373 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
@@ -11,6 +11,9 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirPluginKey
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
@@ -25,7 +28,8 @@ import kotlin.reflect.KProperty
 class FirRegisteredExtension<P : FirExtensionPoint>(
     val extensionsWithAllMode: List<P>,
     val extensionsWithAnnotatedMode: Multimap<AnnotationFqn, P>,
-    val extensionsWithAllInAnnotatedMode: Multimap<AnnotationFqn, P>
+    val extensionsWithAllInAnnotatedMode: Multimap<AnnotationFqn, P>,
+    val extensionsByPluginKey: Multimap<FirPluginKey, P>
 ) {
     fun isEmpty(): Boolean {
         return extensionsWithAllMode.isEmpty() && extensionsWithAnnotatedMode.isEmpty() && extensionsWithAllInAnnotatedMode.isEmpty()
@@ -58,6 +62,7 @@ class FirExtensionPointService(
         val extensionsWithAllMode = mutableListOf<P>()
         val extensionsWithAnnotatedMode = createMultimap<AnnotationFqn, P>()
         val extensionsWithAllInAnnotatedMode = createMultimap<AnnotationFqn, P>()
+        val extensionsByPlugin = createMultimap<FirPluginKey, P>()
         for (extension in extensions) {
             _metaAnnotations += extension.metaAnnotations.keys
             for (metaAnnotation in extension.metaAnnotations.keys) {
@@ -76,13 +81,15 @@ class FirExtensionPointService(
                 }
                 FirExtensionPoint.Mode.ALL -> extensionsWithAllMode += extension
             }
+            extensionsByPlugin.put(extension.key, extension)
         }
         registerComponent(
             extensionClass,
             FirRegisteredExtension(
                 extensionsWithAllMode,
                 extensionsWithAnnotatedMode,
-                extensionsWithAllInAnnotatedMode
+                extensionsWithAllInAnnotatedMode,
+                extensionsByPlugin
             )
         )
     }
@@ -147,6 +154,10 @@ class FirExtensionPointService(
             }
         }
 
+        fun forGeneratedDeclaration(origin: FirDeclarationOrigin.Plugin): Collection<P> {
+            return extensions.extensionsByPluginKey.get(origin.key)
+        }
+
         private fun collectExtensions(result: MutableSet<P>, declaration: FirAnnotationContainer, extensions: Multimap<AnnotationFqn, P>) {
             for (annotation in declaration.annotations) {
                 val fqName = annotation.fqName(session) ?: continue
diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationOrigin.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationOrigin.kt
index e4c0667b680..961c0305ad0 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationOrigin.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/FirDeclarationOrigin.kt
@@ -14,7 +14,7 @@ sealed class FirDeclarationOrigin {
     object FakeOverride : FirDeclarationOrigin()
     object Enhancement : FirDeclarationOrigin()
 
-    class Plugin(val marker: FirPluginMarker) : FirDeclarationOrigin()
+    class Plugin(val key: FirPluginKey) : FirDeclarationOrigin()
 }
 
-abstract class FirPluginMarker
\ No newline at end of file
+abstract class FirPluginKey
\ No newline at end of file