[FIR] Add different modes for extensions and annotations

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

@@ -17,9 +17,17 @@ abstract class FirExtensionPoint(val session: FirSession) {
     abstract val annotations: Set<AnnotationFqn>
     abstract val metaAnnotations: Set<AnnotationFqn>
 
+    abstract val mode: Mode
+
     fun interface Factory<P : FirExtensionPoint> {
         fun create(session: FirSession): P
     }
+
+    enum class Mode {
+        ANNOTATED_ELEMENT,
+        ALL_IN_ANNOTATED_ELEMENT,
+        ALL
+    }
 }
 
 data class FirExtensionPointName(val name: Name) {

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

@@ -10,17 +10,26 @@ 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.expressions.FirAnnotationCall
 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 org.jetbrains.kotlin.name.FqName
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 
-@Suppress("EXPERIMENTAL_FEATURE_WARNING")
-inline class FirRegisteredExtension<P : FirExtensionPoint>(val extensions: Multimap<AnnotationFqn, P>)
+class FirRegisteredExtension<P : FirExtensionPoint>(
+    val extensionsWithAllMode: List<P>,
+    val extensionsWithAnnotatedMode: Multimap<AnnotationFqn, P>,
+    val extensionsWithAllInAnnotatedMode: Multimap<AnnotationFqn, P>
+) {
+    fun isEmpty(): Boolean {
+        return extensionsWithAllMode.isEmpty() && extensionsWithAnnotatedMode.isEmpty() && extensionsWithAllInAnnotatedMode.isEmpty()
+    }
+}
 
 class FirExtensionPointService(
     val session: FirSession
@@ -30,7 +39,7 @@ class FirExtensionPointService(
             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)
+                    return ExtensionsAccessor(thisRef.session, accessor.getValue(thisRef, property))
                 }
             }
         }
@@ -43,15 +52,35 @@ class FirExtensionPointService(
 
     fun <P : FirExtensionPoint> registerExtensions(extensionClass: KClass<P>, extensionFactories: List<FirExtensionPoint.Factory<P>>) {
         val extensions = extensionFactories.map { it.create(session) }
-        val map = createMultimap<AnnotationFqn, P>()
+
+        val extensionsWithAllMode = mutableListOf<P>()
+        val extensionsWithAnnotatedMode = createMultimap<AnnotationFqn, P>()
+        val extensionsWithAllInAnnotatedMode = createMultimap<AnnotationFqn, P>()
         for (extension in extensions) {
             _metaAnnotations += extension.metaAnnotations
-            for (annotation in extension.annotations) {
-                _annotations += annotation
-                map.put(annotation, extension)
+            _annotations += extension.annotations
+            when (extension.mode) {
+                FirExtensionPoint.Mode.ANNOTATED_ELEMENT -> {
+                    for (annotation in extension.annotations) {
+                        extensionsWithAnnotatedMode.put(annotation, extension)
+                    }
+                }
+                FirExtensionPoint.Mode.ALL_IN_ANNOTATED_ELEMENT -> {
+                    for (annotation in extension.annotations) {
+                        extensionsWithAllInAnnotatedMode.put(annotation, extension)
+                    }
+                }
+                FirExtensionPoint.Mode.ALL -> extensionsWithAllMode += extension
             }
         }
-        registerComponent(extensionClass, FirRegisteredExtension(map))
+        registerComponent(
+            extensionClass,
+            FirRegisteredExtension(
+                extensionsWithAllMode,
+                extensionsWithAnnotatedMode,
+                extensionsWithAllInAnnotatedMode
+            )
+        )
     }
 
     val annotations: Set<AnnotationFqn>
@@ -64,17 +93,37 @@ class FirExtensionPointService(
 
     class ExtensionsAccessor<P : FirExtensionPoint>(
         private val session: FirSession,
-        private val extensions: Multimap<AnnotationFqn, P>
+        private val extensions: FirRegisteredExtension<P>
     ) {
-        fun forDeclaration(declaration: FirAnnotationContainer): Collection<P> {
-            if (declaration.annotations.isEmpty()) return emptySet()
-            val result = mutableSetOf<P>()
+        fun forDeclaration(declaration: FirAnnotationContainer, owners: Collection<FirAnnotationContainer>): Collection<P> {
+            if (extensions.isEmpty()) return emptySet()
+            if (
+                (declaration.annotations.isEmpty() || extensions.extensionsWithAnnotatedMode.isEmpty) &&
+                extensions.extensionsWithAllInAnnotatedMode.isEmpty
+            ) {
+                return extensions.extensionsWithAllMode
+            }
+
+            return LinkedHashSet(extensions.extensionsWithAllMode).apply {
+                collectExtensions(this, declaration, extensions.extensionsWithAnnotatedMode)
+                collectExtensions(this, declaration, extensions.extensionsWithAllInAnnotatedMode)
+                for (owner in owners) {
+                    collectExtensions(this, owner, extensions.extensionsWithAllInAnnotatedMode)
+                }
+            }
+        }
+
+        private val FirAnnotationCall.fqName: FqName?
+            get() {
+                val symbol = session.firSymbolProvider.getSymbolByTypeRef<FirRegularClassSymbol>(annotationTypeRef) ?: return null
+                return symbol.classId.asSingleFqName()
+            }
+
+        private fun collectExtensions(result: MutableSet<P>, declaration: FirAnnotationContainer, extensions: Multimap<AnnotationFqn, P>) {
             for (annotation in declaration.annotations) {
-                val symbol = session.firSymbolProvider.getSymbolByTypeRef<FirRegularClassSymbol>(annotation.annotationTypeRef) ?: continue
-                val fqName = symbol.classId.asSingleFqName()
+                val fqName = annotation.fqName ?: continue
                 result += extensions[fqName]
             }
-            return result
         }
     }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirStatusResolveTransformer.kt

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