[FIR] Reimplement FirExtensionService

- Get rid of old FirExtensionsService
- Implement new FirExtensionService
- Remove duplicating code in FirExtensionRegistrar

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt

@@ -45,7 +45,10 @@ import org.jetbrains.kotlin.cli.common.messages.OutputMessageUtil
 import org.jetbrains.kotlin.cli.common.output.writeAll
 import org.jetbrains.kotlin.cli.common.toLogger
 import org.jetbrains.kotlin.cli.jvm.config.*
-import org.jetbrains.kotlin.codegen.*
+import org.jetbrains.kotlin.codegen.ClassBuilderFactories
+import org.jetbrains.kotlin.codegen.CodegenFactory
+import org.jetbrains.kotlin.codegen.DefaultCodegenFactory
+import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.codegen.state.GenerationStateEventCallback
 import org.jetbrains.kotlin.config.*
@@ -54,8 +57,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.backend.Fir2IrConverter
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmClassCodegen
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
-import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.oldExtensionsService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
 import org.jetbrains.kotlin.fir.java.FirLibrarySession
@@ -327,7 +329,7 @@ object KotlinToJVMBytecodeCompiler {
                     dependenciesInfo, provider, librariesScope,
                     project, environment.createPackagePartProvider(librariesScope)
                 )
-                it.oldExtensionsService.registerExtensions(FirExtensionRegistrar.RegisteredExtensions.EMPTY)
+                it.extensionService.registerExtensions(emptyList())
             }
             val firProvider = (session.firProvider as FirProviderImpl)
             val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/extensions/FirAdditionalCheckersExtension.kt

@@ -8,9 +8,9 @@ package org.jetbrains.kotlin.fir.analysis.extensions
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
-import org.jetbrains.kotlin.fir.extensions.*
+import org.jetbrains.kotlin.fir.extensions.AbstractFirAdditionalCheckersExtension
-import kotlin.properties.ReadOnlyProperty
+import org.jetbrains.kotlin.fir.extensions.FirExtensionPointName
-import kotlin.reflect.KProperty
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 
 abstract class FirAdditionalCheckersExtension(session: FirSession) : AbstractFirAdditionalCheckersExtension(session) {
     companion object {
@@ -22,31 +22,8 @@ abstract class FirAdditionalCheckersExtension(session: FirSession) : AbstractFir
 
     final override val name: FirExtensionPointName
         get() = NAME
-
-    final override val directlyApplicableAnnotations: Set<AnnotationFqn>
-        get() = emptySet()
-
-    final override val childrenApplicableAnnotations: Set<AnnotationFqn>
-        get() = emptySet()
-
-    final override val metaAnnotations: Map<AnnotationFqn, MetaAnnotationMode>
-        get() = emptyMap()
-
-    final override val mode: Mode
-        get() = Mode.ALL
 }
 
-val FirOldExtensionsService.additionalCheckers: List<FirAdditionalCheckersExtension> by FirOldExtensionsService.additionalCheckers()
+val FirExtensionService.additionalCheckers: List<FirAdditionalCheckersExtension> by FirExtensionService.registeredExtensions(
-
+    AbstractFirAdditionalCheckersExtension::class
-private fun FirOldExtensionsService.Companion.additionalCheckers(): ReadOnlyProperty<FirOldExtensionsService, List<FirAdditionalCheckersExtension>> {
+)
-    val accessor = generateAccessor<FirRegisteredExtension<AbstractFirAdditionalCheckersExtension>, AbstractFirAdditionalCheckersExtension>(
-        AbstractFirAdditionalCheckersExtension::class
-    )
-    return object : ReadOnlyProperty<FirOldExtensionsService, List<FirAdditionalCheckersExtension>> {
-        override fun getValue(thisRef: FirOldExtensionsService, property: KProperty<*>): List<FirAdditionalCheckersExtension> {
-            val extensions = accessor.getValue(thisRef, property)
-            @Suppress("UNCHECKED_CAST")
-            return extensions.extensionsWithAllMode as List<FirAdditionalCheckersExtension>
-        }
-    }
-}

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt

@@ -15,7 +15,8 @@ import org.jetbrains.kotlin.fir.FirSession
 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.FirExtensionService
+import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
 import org.jetbrains.kotlin.fir.resolve.FirProvider
@@ -123,8 +124,13 @@ class FirLibrarySession private constructor(
         )
 
         registerComponent(
-            FirOldExtensionsService::class,
+            FirExtensionService::class,
-            FirOldExtensionsService(this)
+            FirExtensionService(this)
+        )
+
+        registerComponent(
+            FirRegisteredPluginAnnotations::class,
+            FirRegisteredPluginAnnotations.create(this)
         )
 
         registerComponent(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirModuleBasedSession.kt

@@ -6,10 +6,13 @@
 package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.analyzer.ModuleInfo
-import org.jetbrains.kotlin.fir.extensions.FirOldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
+import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
+import org.jetbrains.kotlin.fir.resolve.FirQualifierResolver
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
-import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.FirTypeResolver
-import org.jetbrains.kotlin.fir.resolve.impl.*
+import org.jetbrains.kotlin.fir.resolve.impl.FirQualifierResolverImpl
+import org.jetbrains.kotlin.fir.resolve.impl.FirTypeResolverImpl
 import org.jetbrains.kotlin.fir.scopes.impl.FirDeclaredMemberScopeProvider
 
 abstract class FirModuleBasedSession(override val moduleInfo: ModuleInfo, sessionProvider: FirSessionProvider?) :
@@ -18,7 +21,8 @@ abstract class FirModuleBasedSession(override val moduleInfo: ModuleInfo, sessio
         registerComponent(FirQualifierResolver::class, FirQualifierResolverImpl(this))
         registerComponent(FirTypeResolver::class, FirTypeResolverImpl(this))
         registerComponent(FirDeclaredMemberScopeProvider::class, FirDeclaredMemberScopeProvider())
-        registerComponent(FirOldExtensionsService::class, FirOldExtensionsService(this))
+        registerComponent(FirExtensionService::class, FirExtensionService(this))
+        registerComponent(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.create(this))
         registerComponent(FirPredicateBasedProvider::class, FirPredicateBasedProvider.create(this))
     }
 }

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

@@ -17,7 +17,7 @@ import kotlin.reflect.KClass
  *  - add registration of new files
  *  - check that annotations or meta-annotations is not empty
  */
-abstract class FirClassGenerationExtension(session: FirSession) : FirExtension(session) {
+abstract class FirClassGenerationExtension(session: FirSession) : FirPredicateBasedExtension(session) {
     companion object {
         val NAME = FirExtensionPointName("StatusTransformer")
     }
@@ -25,8 +25,6 @@ abstract class FirClassGenerationExtension(session: FirSession) : FirExtension(s
     final override val name: FirExtensionPointName
         get() = NAME
 
-    final override val mode: Mode = Mode.ANNOTATED_ELEMENT
-
     final override val extensionType: KClass<out FirExtension> = FirClassGenerationExtension::class
 
     abstract fun <T> generateClass(
@@ -43,4 +41,4 @@ abstract class FirClassGenerationExtension(session: FirSession) : FirExtension(s
     fun interface Factory : FirExtension.Factory<FirClassGenerationExtension>
 }
 
-val FirOldExtensionsService.classGenerationExtensions: FirOldExtensionsService.ExtensionsAccessor<FirClassGenerationExtension> by FirOldExtensionsService.registeredExtensions()
+val FirExtensionService.classGenerationExtensions: List<FirClassGenerationExtension> by FirExtensionService.registeredExtensions()

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

@@ -7,20 +7,23 @@ package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirPluginKey
+import org.jetbrains.kotlin.fir.extensions.predicate.DeclarationPredicate
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import kotlin.reflect.KClass
 
 typealias AnnotationFqn = FqName
 
+/*
+ * Accessing extensions
+ *
+ * - specific extension -> all extension instances: extension accessor on `FirExtensionService`
+ * - all declarations matching extension with predicate -> FirPredicateBasedProvider.getSymbolsByPredicate
+ * - all specific extensions interested in specific declaration -> TODO with StateMachine
+ */
 abstract class FirExtension(val session: FirSession) {
     abstract val name: FirExtensionPointName
 
-    abstract val directlyApplicableAnnotations: Set<AnnotationFqn>
-    abstract val childrenApplicableAnnotations: Set<AnnotationFqn>
-    abstract val metaAnnotations: Map<AnnotationFqn, MetaAnnotationMode>
-
-    abstract val mode: Mode
     abstract val key: FirPluginKey
 
     internal abstract val extensionType: KClass<out FirExtension>
@@ -28,17 +31,10 @@ abstract class FirExtension(val session: FirSession) {
     fun interface Factory<out P : FirExtension> {
         fun create(session: FirSession): P
     }
+}
 
-    enum class Mode {
+abstract class FirPredicateBasedExtension(session: FirSession) : FirExtension(session) {
-        ANNOTATED_ELEMENT,
+    abstract val predicate: DeclarationPredicate
-        ALL
-    }
-
-    enum class MetaAnnotationMode(val directed: Boolean, val children: Boolean) {
-        ANNOTATED_DECLARATION(directed = true, children = false),
-        CHILDREN_DECLARATION(directed = false, children = true),
-        ANNOTATED_AND_CHILDREN(directed = true, children = true)
-    }
 }
 
 data class FirExtensionPointName(val name: Name) {

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
 import org.jetbrains.kotlin.fir.FirSession
+import kotlin.reflect.KClass
 
 abstract class FirExtensionRegistrar {
     companion object : ProjectExtensionDescriptor<FirExtensionRegistrar>(
@@ -14,51 +15,48 @@ abstract class FirExtensionRegistrar {
         FirExtensionRegistrar::class.java
     )
 
+    protected abstract fun ExtensionRegistrarContext.configurePlugin()
+
     protected inner class ExtensionRegistrarContext {
         @JvmName("plusStatusTransformerExtension")
         operator fun ((FirSession) -> FirStatusTransformerExtension).unaryPlus() {
-            statusTransformerExtensions += FirStatusTransformerExtension.Factory { this.invoke(it) }
+            registerExtension(FirStatusTransformerExtension::class, FirStatusTransformerExtension.Factory { this.invoke(it) })
         }
 
         @JvmName("plusClassGenerationExtension")
         operator fun ((FirSession) -> FirClassGenerationExtension).unaryPlus() {
-            classGenerationExtensions += FirClassGenerationExtension.Factory { this.invoke(it) }
+            registerExtension(FirClassGenerationExtension::class, FirClassGenerationExtension.Factory { this.invoke(it) })
         }
 
         @JvmName("plusAdditionalCheckersExtension")
         operator fun ((FirSession) -> AbstractFirAdditionalCheckersExtension).unaryPlus() {
-            additionalCheckersExtensions += AbstractFirAdditionalCheckersExtension.Factory { this.invoke(it) }
+            registerExtension(
+                AbstractFirAdditionalCheckersExtension::class,
+                AbstractFirAdditionalCheckersExtension.Factory { this.invoke(it) }
+            )
         }
     }
 
-    protected abstract fun ExtensionRegistrarContext.configurePlugin()
+    fun configure(): Collection<RegisteredExtensionsFactories<*>> {
-
-    fun configure(): RegisteredExtensions {
         ExtensionRegistrarContext().configurePlugin()
-        return RegisteredExtensions(
+        return map.values
-            statusTransformerExtensions,
-            classGenerationExtensions,
-            additionalCheckersExtensions,
-        )
     }
 
-    private val statusTransformerExtensions: MutableList<FirStatusTransformerExtension.Factory> = mutableListOf()
+    class RegisteredExtensionsFactories<P : FirExtension>(val kClass: KClass<P>) {
-    private val classGenerationExtensions: MutableList<FirClassGenerationExtension.Factory> = mutableListOf()
+        val extensionFactories: MutableList<FirExtension.Factory<P>> = mutableListOf()
-    private val additionalCheckersExtensions: MutableList<AbstractFirAdditionalCheckersExtension.Factory> = mutableListOf()
+    }
 
-    class RegisteredExtensions(
+    private val map: MutableMap<KClass<out FirExtension>, RegisteredExtensionsFactories<*>> = mutableMapOf()
-        val statusTransformerExtensions: List<FirStatusTransformerExtension.Factory>,
+
-        val classGenerationExtensions: List<FirClassGenerationExtension.Factory>,
+    private fun <P : FirExtension> registerExtension(kClass: KClass<out P>, factory: FirExtension.Factory<P>) {
-        val additionalCheckersExtensions: List<AbstractFirAdditionalCheckersExtension.Factory>,
+        @Suppress("UNCHECKED_CAST")
-    ) {
+        val registeredExtensions = map.computeIfAbsent(kClass) { RegisteredExtensionsFactories(kClass) } as RegisteredExtensionsFactories<P>
-        companion object {
+        registeredExtensions.extensionFactories += factory
-            val EMPTY = RegisteredExtensions(emptyList(), emptyList(), emptyList())
-        }
     }
 }
 
-fun FirOldExtensionsService.registerExtensions(extensions: FirExtensionRegistrar.RegisteredExtensions) {
+@OptIn(PluginServicesInitialization::class)
-    registerExtensions(FirStatusTransformerExtension::class, extensions.statusTransformerExtensions)
+fun FirExtensionService.registerExtensions(registeredExtensions: Collection<FirExtensionRegistrar.RegisteredExtensionsFactories<*>>) {
-    registerExtensions(FirClassGenerationExtension::class, extensions.classGenerationExtensions)
+    registeredExtensions.forEach { registerExtensions(it.kClass, it.extensionFactories) }
-    registerExtensions(AbstractFirAdditionalCheckersExtension::class, extensions.additionalCheckersExtensions)
+    session.registeredPluginAnnotations.initialize()
-}
+}

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

@@ -0,0 +1,53 @@
+/*
+ * 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 org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.utils.ArrayMapAccessor
+import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner
+import org.jetbrains.kotlin.fir.utils.TypeRegistry
+import kotlin.reflect.KClass
+
+@RequiresOptIn
+annotation class PluginServicesInitialization
+
+class FirExtensionService(val session: FirSession) : ComponentArrayOwner<FirExtension, List<FirExtension>>(), FirSessionComponent {
+    companion object : TypeRegistry<FirExtension, List<FirExtension>>() {
+        inline fun <reified P : FirExtension, V : List<P>> registeredExtensions(): ArrayMapAccessor<FirExtension, List<FirExtension>, V> {
+            return generateAccessor(P::class)
+        }
+
+        fun <P : FirExtension, V : List<P>> registeredExtensions(kClass: KClass<P>): ArrayMapAccessor<FirExtension, List<FirExtension>, V> {
+            return generateAccessor(kClass)
+        }
+    }
+
+    override val typeRegistry: TypeRegistry<FirExtension, List<FirExtension>>
+        get() = Companion
+
+    var registeredExtensionsSize: Int = 0
+        private set
+
+    @PluginServicesInitialization
+    fun registerExtensions(extensionClass: KClass<out FirExtension>, extensionFactories: List<FirExtension.Factory<*>>) {
+        registeredExtensionsSize += extensionFactories.size
+        registerComponent(
+            extensionClass,
+            extensionFactories.map { it.create(session) }
+        )
+    }
+
+    @PluginServicesInitialization
+    fun getAllExtensions(): List<FirExtension> {
+        return arrayMap.flatten()
+    }
+}
+
+val FirSession.extensionService: FirExtensionService by FirSession.sessionComponentAccessor()
+
+val FirExtensionService.hasExtensions: Boolean
+    get() = registeredExtensionsSize > 0

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

@@ -7,24 +7,20 @@ 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.FirAnnotationContainer
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
 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.fqName
-import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
+import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner
-import org.jetbrains.kotlin.fir.resolve.getSymbolByTypeRef
+import org.jetbrains.kotlin.fir.utils.TypeRegistry
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.utils.*
-import org.jetbrains.kotlin.name.FqName
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 
-class FirRegisteredExtension<P : FirExtension>(
+class OldFirRegisteredExtension<P : FirExtension>(
     val extensionsWithAllMode: List<P>,
     val extensionsWithAnnotatedMode: Multimap<AnnotationFqn, P>,
     val extensionsWithAllInAnnotatedMode: Multimap<AnnotationFqn, P>,
@@ -38,9 +34,9 @@ class FirRegisteredExtension<P : FirExtension>(
 @Deprecated("")
 class FirOldExtensionsService(
     val session: FirSession
-) : ComponentArrayOwner<FirExtension, FirRegisteredExtension<*>>(), FirSessionComponent {
+) : ComponentArrayOwner<FirExtension, OldFirRegisteredExtension<*>>(), FirSessionComponent {
-    companion object : TypeRegistry<FirExtension, FirRegisteredExtension<*>>() {
+    companion object : TypeRegistry<FirExtension, OldFirRegisteredExtension<*>>() {
-        inline fun <reified P : FirExtension, V : FirRegisteredExtension<P>> registeredExtensions(): ReadOnlyProperty<FirOldExtensionsService, ExtensionsAccessor<P>> {
+        inline fun <reified P : FirExtension, V : OldFirRegisteredExtension<P>> registeredExtensions(): ReadOnlyProperty<FirOldExtensionsService, ExtensionsAccessor<P>> {
             val accessor = generateAccessor<V, P>(P::class)
             return object : ReadOnlyProperty<FirOldExtensionsService, ExtensionsAccessor<P>> {
                 override fun getValue(thisRef: FirOldExtensionsService, property: KProperty<*>): ExtensionsAccessor<P> {
@@ -52,73 +48,73 @@ class FirOldExtensionsService(
         private fun <K, V> createMultimap(): Multimap<K, V> = LinkedHashMultimap.create()
     }
 
-    override val typeRegistry: TypeRegistry<FirExtension, FirRegisteredExtension<*>>
+    override val typeRegistry: TypeRegistry<FirExtension, OldFirRegisteredExtension<*>>
         get() = Companion
 
     fun <P : FirExtension> registerExtensions(extensionClass: KClass<P>, extensionFactories: List<FirExtension.Factory<P>>) {
-        registeredExtensionsSize += extensionFactories.size
+//        registeredExtensionsSize += extensionFactories.size
-        val extensions = extensionFactories.map { it.create(session) }
+//        val extensions = extensionFactories.map { it.create(session) }
-
+//
-        val extensionsWithAllMode = mutableListOf<P>()
+//        val extensionsWithAllMode = mutableListOf<P>()
-        val extensionsWithAnnotatedMode = createMultimap<AnnotationFqn, P>()
+//        val extensionsWithAnnotatedMode = createMultimap<AnnotationFqn, P>()
-        val extensionsWithAllInAnnotatedMode = createMultimap<AnnotationFqn, P>()
+//        val extensionsWithAllInAnnotatedMode = createMultimap<AnnotationFqn, P>()
-        val extensionsByPlugin = createMultimap<FirPluginKey, P>()
+//        val extensionsByPlugin = createMultimap<FirPluginKey, P>()
-        for (extension in extensions) {
+//        for (extension in extensions) {
-            _metaAnnotations += extension.metaAnnotations.keys
+//            _metaAnnotations += extension.metaAnnotations.keys
-            for (metaAnnotation in extension.metaAnnotations.keys) {
+//            for (metaAnnotation in extension.metaAnnotations.keys) {
-                extensionsWithMetaAnnotations.put(metaAnnotation, extension)
+//                extensionsWithMetaAnnotations.put(metaAnnotation, extension)
-            }
+//            }
-            _annotations += extension.directlyApplicableAnnotations
+//            _annotations += extension.directlyApplicableAnnotations
-            _annotations += extension.childrenApplicableAnnotations
+//            _annotations += extension.childrenApplicableAnnotations
-            when (extension.mode) {
+//            when (extension.mode) {
-                FirExtension.Mode.ANNOTATED_ELEMENT -> {
+//                FirExtension.Mode.ANNOTATED_ELEMENT -> {
-                    for (annotation in extension.directlyApplicableAnnotations) {
+//                    for (annotation in extension.directlyApplicableAnnotations) {
-                        extensionsWithAnnotatedMode.put(annotation, extension)
+//                        extensionsWithAnnotatedMode.put(annotation, extension)
-                    }
+//                    }
-                    for (annotation in extension.childrenApplicableAnnotations) {
+//                    for (annotation in extension.childrenApplicableAnnotations) {
-                        extensionsWithAllInAnnotatedMode.put(annotation, extension)
+//                        extensionsWithAllInAnnotatedMode.put(annotation, extension)
-                    }
+//                    }
-                }
+//                }
-                FirExtension.Mode.ALL -> extensionsWithAllMode += extension
+//                FirExtension.Mode.ALL -> extensionsWithAllMode += extension
-            }
+//            }
-            extensionsByPlugin.put(extension.key, extension)
+//            extensionsByPlugin.put(extension.key, extension)
-        }
+//        }
-        registerComponent(
+//        registerComponent(
-            extensionClass,
+//            extensionClass,
-            FirRegisteredExtension(
+//            FirRegisteredExtension(
-                extensionsWithAllMode,
+//                extensionsWithAllMode,
-                extensionsWithAnnotatedMode,
+//                extensionsWithAnnotatedMode,
-                extensionsWithAllInAnnotatedMode,
+//                extensionsWithAllInAnnotatedMode,
-                extensionsByPlugin
+//                extensionsByPlugin
-            )
+//            )
-        )
+//        )
     }
 
     @Deprecated("")
     fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotations: Collection<FirRegularClass>) {
-        for (annotation in annotations) {
+//        for (annotation in annotations) {
-            require(annotation.classKind == ClassKind.ANNOTATION_CLASS)
+//            require(annotation.classKind == ClassKind.ANNOTATION_CLASS)
-            val fqName = annotation.symbol.classId.asSingleFqName()
+//            val fqName = annotation.symbol.classId.asSingleFqName()
-            _annotations += fqName
+//            _annotations += fqName
-            userDefinedAnnotations.put(metaAnnotation, fqName)
+//            userDefinedAnnotations.put(metaAnnotation, fqName)
-            val extensions = extensionsWithMetaAnnotations[metaAnnotation]
+//            val extensions = extensionsWithMetaAnnotations[metaAnnotation]
-            if (extensions.isEmpty()) return
+//            if (extensions.isEmpty()) return
-            for (extension in extensions) {
+//            for (extension in extensions) {
-                val registeredExtensions = this[extension.extensionType]
+//                val registeredExtensions = this[extension.extensionType]
-
+//
-                val metaAnnotationMode = extension.metaAnnotations.getValue(metaAnnotation)
+//                val metaAnnotationMode = extension.metaAnnotations.getValue(metaAnnotation)
-
+//
-                if (metaAnnotationMode.directed) {
+//                if (metaAnnotationMode.directed) {
-                    @Suppress("UNCHECKED_CAST")
+//                    @Suppress("UNCHECKED_CAST")
-                    (registeredExtensions.extensionsWithAnnotatedMode as Multimap<AnnotationFqn, FirExtension>).put(fqName, extension)
+//                    (registeredExtensions.extensionsWithAnnotatedMode as Multimap<AnnotationFqn, FirExtension>).put(fqName, extension)
-                }
+//                }
-
+//
-                if (metaAnnotationMode.children) {
+//                if (metaAnnotationMode.children) {
-                    @Suppress("UNCHECKED_CAST")
+//                    @Suppress("UNCHECKED_CAST")
-                    (registeredExtensions.extensionsWithAllInAnnotatedMode as Multimap<AnnotationFqn, FirExtension>).put(fqName, extension)
+//                    (registeredExtensions.extensionsWithAllInAnnotatedMode as Multimap<AnnotationFqn, FirExtension>).put(fqName, extension)
-                }
+//                }
-            }
+//            }
-        }
+//        }
     }
 
     @Deprecated("")
@@ -142,7 +138,7 @@ class FirOldExtensionsService(
 
     class ExtensionsAccessor<P : FirExtension>(
         private val session: FirSession,
-        private val extensions: FirRegisteredExtension<P>
+        private val extensions: OldFirRegisteredExtension<P>
     ) {
         fun forDeclaration(declaration: FirAnnotationContainer, owners: Collection<FirAnnotationContainer>): Collection<P> {
             if (extensions.isEmpty()) return emptySet()
@@ -178,10 +174,5 @@ class FirOldExtensionsService(
 @Deprecated("")
 val FirSession.oldExtensionsService: FirOldExtensionsService by FirSession.sessionComponentAccessor()
 
-fun FirAnnotationCall.fqName(session: FirSession): FqName? {
-    val symbol = session.firSymbolProvider.getSymbolByTypeRef<FirRegularClassSymbol>(annotationTypeRef) ?: return null
-    return symbol.classId.asSingleFqName()
-}
-
 val FirOldExtensionsService.hasExtensions: Boolean
     get() = registeredExtensionsSize > 0

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

@@ -13,10 +13,10 @@ 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 {
+abstract class FirRegisteredPluginAnnotations(val session: FirSession) : FirSessionComponent {
     companion object {
-        fun create(): FirRegisteredPluginAnnotations {
+        fun create(session: FirSession): FirRegisteredPluginAnnotations {
-            return FirRegisteredPluginAnnotationsImpl()
+            return FirRegisteredPluginAnnotationsImpl(session)
         }
     }
 
@@ -24,14 +24,15 @@ abstract class FirRegisteredPluginAnnotations : FirSessionComponent {
     abstract val metaAnnotations: Set<AnnotationFqn>
     abstract fun getAnnotationsWithMetaAnnotation(metaAnnotation: AnnotationFqn): Collection<AnnotationFqn>
 
-    abstract fun registerAnnotations(annotations: Collection<AnnotationFqn>)
-    abstract fun registerMetaAnnotations(metaAnnotations: Collection<AnnotationFqn>)
     abstract fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection<FirRegularClass>)
 
     abstract fun getAnnotationsForPredicate(predicate: DeclarationPredicate): Set<AnnotationFqn>
+
+    @PluginServicesInitialization
+    abstract fun initialize()
 }
 
-private class FirRegisteredPluginAnnotationsImpl : FirRegisteredPluginAnnotations() {
+private class FirRegisteredPluginAnnotationsImpl(session: FirSession) : FirRegisteredPluginAnnotations(session) {
     override val annotations: MutableSet<AnnotationFqn> = mutableSetOf()
     override val metaAnnotations: MutableSet<AnnotationFqn> = mutableSetOf()
 
@@ -44,18 +45,10 @@ private class FirRegisteredPluginAnnotationsImpl : FirRegisteredPluginAnnotation
         return userDefinedAnnotations[metaAnnotation]
     }
 
-    override fun registerAnnotations(annotations: Collection<AnnotationFqn>) {
-        this.annotations += annotations
-    }
-
-    override fun registerMetaAnnotations(metaAnnotations: Collection<AnnotationFqn>) {
-        this.metaAnnotations += metaAnnotations
-    }
-
     override fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotationClasses: Collection<FirRegularClass>) {
         require(annotationClasses.all { it.classKind == ClassKind.ANNOTATION_CLASS })
         val annotations = annotationClasses.map { it.symbol.classId.asSingleFqName() }
-        registerAnnotations(annotations)
+        this.annotations += annotations
         userDefinedAnnotations.putAll(metaAnnotation, annotations)
     }
 
@@ -70,6 +63,16 @@ private class FirRegisteredPluginAnnotationsImpl : FirRegisteredPluginAnnotation
         result += predicate.annotations
         return result
     }
+
+    @PluginServicesInitialization
+    override fun initialize() {
+        for (extension in session.extensionService.getAllExtensions()) {
+            if (extension !is FirPredicateBasedExtension) continue
+            val predicate = extension.predicate
+            annotations += predicate.annotations
+            metaAnnotations += predicate.metaAnnotations
+        }
+    }
 }
 
 val FirSession.registeredPluginAnnotations: FirRegisteredPluginAnnotations by FirSession.sessionComponentAccessor()

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

@@ -13,7 +13,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclarationStatus
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import kotlin.reflect.KClass
 
-abstract class FirStatusTransformerExtension(session: FirSession) : FirExtension(session) {
+abstract class FirStatusTransformerExtension(session: FirSession) : FirPredicateBasedExtension(session) {
     companion object {
         val NAME = FirExtensionPointName("StatusTransformer")
     }
@@ -28,7 +28,7 @@ abstract class FirStatusTransformerExtension(session: FirSession) : FirExtension
     fun interface Factory : FirExtension.Factory<FirStatusTransformerExtension>
 }
 
-val FirOldExtensionsService.statusTransformerExtensions: FirOldExtensionsService.ExtensionsAccessor<FirStatusTransformerExtension> by FirOldExtensionsService.registeredExtensions()
+val FirExtensionService.statusTransformerExtensions: List<FirStatusTransformerExtension> by FirExtensionService.registeredExtensions()
 
 inline fun FirDeclarationStatus.transform(
     visibility: Visibility = this.visibility,

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -11,10 +11,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.diagnostics.ConeStubDiagnostic
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.fir.expressions.FirResolvable
-import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.expressions.builder.FirResolvedReifiedParameterReferenceBuilder
 import org.jetbrains.kotlin.fir.expressions.builder.buildExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.builder.buildResolvedQualifier
@@ -38,6 +35,7 @@ import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.impl.ConeClassLikeTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
@@ -433,3 +431,8 @@ fun CallableId.isIteratorHasNext() =
 
 fun CallableId.isIterator() =
     callableName.asString() == "iterator" && packageName.asString() == "kotlin.collections"
+
+fun FirAnnotationCall.fqName(session: FirSession): FqName? {
+    val symbol = session.firSymbolProvider.getSymbolByTypeRef<FirRegularClassSymbol>(annotationTypeRef) ?: return null
+    return symbol.classId.asSingleFqName()
+}

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirBaseDiagnosticsTest.kt

@@ -26,9 +26,8 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.extensions.FirOldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
-import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.oldExtensionsService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
 import org.jetbrains.kotlin.fir.java.FirLibrarySession
@@ -80,7 +79,7 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
             builtInsModuleInfo, sessionProvider, allProjectScope, project,
             environment.createPackagePartProvider(allProjectScope)
         ).also {
-            registerFirExtensions(it.oldExtensionsService)
+            registerFirExtensions(it.extensionService)
         }
 
         val configToSession = modules.mapValues { (config, info) ->
@@ -89,7 +88,7 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
                 project,
                 moduleFiles.mapNotNull { it.ktFile })
             FirJavaModuleBasedSession(info, sessionProvider, scope).also {
-                registerFirExtensions(it.oldExtensionsService)
+                registerFirExtensions(it.extensionService)
             }
         }
 
@@ -109,8 +108,8 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         runAnalysis(testDataFile, files, firFilesPerSession)
     }
 
-    open fun registerFirExtensions(service: FirOldExtensionsService) {
+    open fun registerFirExtensions(service: FirExtensionService) {
-        service.registerExtensions(FirExtensionRegistrar.RegisteredExtensions.EMPTY)
+        service.registerExtensions(emptyList())
     }
 
     private fun mapKtFilesToFirFiles(session: FirSession, ktFiles: List<KtFile>, firFiles: MutableList<FirFile>, useLightTree: Boolean) {

compiler/tests-common/tests/org/jetbrains/kotlin/fir/SessionTestUtils.kt

@@ -8,8 +8,7 @@ package org.jetbrains.kotlin.fir
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
-import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.oldExtensionsService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
 import org.jetbrains.kotlin.fir.java.FirLibrarySession
@@ -35,7 +34,7 @@ fun createSession(
     val provider = FirProjectSessionProvider(project)
     return FirJavaModuleBasedSession(moduleInfo, provider, sourceScope).also {
         createSessionForDependencies(project, provider, moduleInfo, librariesScope, packagePartProvider)
-        it.oldExtensionsService.registerExtensions(FirExtensionRegistrar.RegisteredExtensions.EMPTY)
+        it.extensionService.registerExtensions(emptyList())
     }
 }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirModuleResolveState.kt

@@ -14,8 +14,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.extensions.extensionService
-import org.jetbrains.kotlin.fir.extensions.oldExtensionsService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
@@ -40,7 +39,7 @@ interface FirModuleResolveState {
                 sessionProvider.sessionCache[moduleInfo] = moduleBasedSession
             }
             session.also {
-                it.oldExtensionsService.registerExtensions(FirExtensionRegistrar.RegisteredExtensions.EMPTY)
+                it.extensionService.registerExtensions(emptyList())
             }
         }
     }

idea/idea-fir/tests/org/jetbrains/kotlin/idea/fir/AbstractFirMultiModuleResolveTest.kt

@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.dependenciesWithoutSelf
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.extensions.oldExtensionsService
-import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
 import org.jetbrains.kotlin.fir.java.FirLibrarySession
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
@@ -72,7 +71,7 @@ abstract class AbstractFirMultiModuleResolveTest : AbstractMultiModuleTest() {
     private fun createSession(module: Module, provider: FirProjectSessionProvider): FirJavaModuleBasedSession {
         val moduleInfo = module.productionSourceInfo()!!
         return FirJavaModuleBasedSession(moduleInfo, provider, moduleInfo.contentScope()).also {
-            it.oldExtensionsService.registerExtensions(FirExtensionRegistrar.RegisteredExtensions.EMPTY)
+            it.oldExtensionsService.registerExtensions(FirExtensionRegistrar.OldRegisteredExtensions.EMPTY)
         }
     }