[FIR] Reimplement FirExtensionService

- Get rid of old FirExtensionsService - Implement new FirExtensionService - Remove duplicating code in FirExtensionRegistrar
2020-05-20 13:08:32 +03:00
parent 846db641be
commit 594a854b4e
16 changed files with 226 additions and 199 deletions
@@ -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.oldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.extensionService
 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)
@@ -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 kotlin.properties.ReadOnlyProperty
-import kotlin.reflect.KProperty
+import org.jetbrains.kotlin.fir.extensions.AbstractFirAdditionalCheckersExtension
+import org.jetbrains.kotlin.fir.extensions.FirExtensionPointName
+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()
-
-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>
-        }
-    }
-}
+val FirExtensionService.additionalCheckers: List<FirAdditionalCheckersExtension> by FirExtensionService.registeredExtensions(
+    AbstractFirAdditionalCheckersExtension::class
+)
@@ -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,
-            FirOldExtensionsService(this)
+            FirExtensionService::class,
+            FirExtensionService(this)
+        )
+
+        registerComponent(
+            FirRegisteredPluginAnnotations::class,
+            FirRegisteredPluginAnnotations.create(this)
        )

        registerComponent(
@@ -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.impl.*
+import org.jetbrains.kotlin.fir.resolve.FirTypeResolver
+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))
    }
 }
@@ -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()
@@ -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 {
-        ANNOTATED_ELEMENT,
-        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)
-    }
+abstract class FirPredicateBasedExtension(session: FirSession) : FirExtension(session) {
+    abstract val predicate: DeclarationPredicate
 }

 data class FirExtensionPointName(val name: Name) {
@@ -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(): RegisteredExtensions {
+    fun configure(): Collection<RegisteredExtensionsFactories<*>> {
        ExtensionRegistrarContext().configurePlugin()
-        return RegisteredExtensions(
-            statusTransformerExtensions,
-            classGenerationExtensions,
-            additionalCheckersExtensions,
-        )
+        return map.values
    }

-    private val statusTransformerExtensions: MutableList<FirStatusTransformerExtension.Factory> = mutableListOf()
-    private val classGenerationExtensions: MutableList<FirClassGenerationExtension.Factory> = mutableListOf()
-    private val additionalCheckersExtensions: MutableList<AbstractFirAdditionalCheckersExtension.Factory> = mutableListOf()
+    class RegisteredExtensionsFactories<P : FirExtension>(val kClass: KClass<P>) {
+        val extensionFactories: MutableList<FirExtension.Factory<P>> = mutableListOf()
+    }

-    class RegisteredExtensions(
-        val statusTransformerExtensions: List<FirStatusTransformerExtension.Factory>,
-        val classGenerationExtensions: List<FirClassGenerationExtension.Factory>,
-        val additionalCheckersExtensions: List<AbstractFirAdditionalCheckersExtension.Factory>,
-    ) {
-        companion object {
-            val EMPTY = RegisteredExtensions(emptyList(), emptyList(), emptyList())
-        }
+    private val map: MutableMap<KClass<out FirExtension>, RegisteredExtensionsFactories<*>> = mutableMapOf()
+
+    private fun <P : FirExtension> registerExtension(kClass: KClass<out P>, factory: FirExtension.Factory<P>) {
+        @Suppress("UNCHECKED_CAST")
+        val registeredExtensions = map.computeIfAbsent(kClass) { RegisteredExtensionsFactories(kClass) } as RegisteredExtensionsFactories<P>
+        registeredExtensions.extensionFactories += factory
    }
 }

-fun FirOldExtensionsService.registerExtensions(extensions: FirExtensionRegistrar.RegisteredExtensions) {
-    registerExtensions(FirStatusTransformerExtension::class, extensions.statusTransformerExtensions)
-    registerExtensions(FirClassGenerationExtension::class, extensions.classGenerationExtensions)
-    registerExtensions(AbstractFirAdditionalCheckersExtension::class, extensions.additionalCheckersExtensions)
-}
+@OptIn(PluginServicesInitialization::class)
+fun FirExtensionService.registerExtensions(registeredExtensions: Collection<FirExtensionRegistrar.RegisteredExtensionsFactories<*>>) {
+    registeredExtensions.forEach { registerExtensions(it.kClass, it.extensionFactories) }
+    session.registeredPluginAnnotations.initialize()
+}
@@ -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
@@ -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.firSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.getSymbolByTypeRef
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.utils.*
-import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.fir.resolve.fqName
+import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner
+import org.jetbrains.kotlin.fir.utils.TypeRegistry
 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 {
-    companion object : TypeRegistry<FirExtension, FirRegisteredExtension<*>>() {
-        inline fun <reified P : FirExtension, V : FirRegisteredExtension<P>> registeredExtensions(): ReadOnlyProperty<FirOldExtensionsService, ExtensionsAccessor<P>> {
+) : ComponentArrayOwner<FirExtension, OldFirRegisteredExtension<*>>(), FirSessionComponent {
+    companion object : TypeRegistry<FirExtension, OldFirRegisteredExtension<*>>() {
+        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
-        val extensions = extensionFactories.map { it.create(session) }
-
-        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) {
-                extensionsWithMetaAnnotations.put(metaAnnotation, extension)
-            }
-            _annotations += extension.directlyApplicableAnnotations
-            _annotations += extension.childrenApplicableAnnotations
-            when (extension.mode) {
-                FirExtension.Mode.ANNOTATED_ELEMENT -> {
-                    for (annotation in extension.directlyApplicableAnnotations) {
-                        extensionsWithAnnotatedMode.put(annotation, extension)
-                    }
-                    for (annotation in extension.childrenApplicableAnnotations) {
-                        extensionsWithAllInAnnotatedMode.put(annotation, extension)
-                    }
-                }
-                FirExtension.Mode.ALL -> extensionsWithAllMode += extension
-            }
-            extensionsByPlugin.put(extension.key, extension)
-        }
-        registerComponent(
-            extensionClass,
-            FirRegisteredExtension(
-                extensionsWithAllMode,
-                extensionsWithAnnotatedMode,
-                extensionsWithAllInAnnotatedMode,
-                extensionsByPlugin
-            )
-        )
+//        registeredExtensionsSize += extensionFactories.size
+//        val extensions = extensionFactories.map { it.create(session) }
+//
+//        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) {
+//                extensionsWithMetaAnnotations.put(metaAnnotation, extension)
+//            }
+//            _annotations += extension.directlyApplicableAnnotations
+//            _annotations += extension.childrenApplicableAnnotations
+//            when (extension.mode) {
+//                FirExtension.Mode.ANNOTATED_ELEMENT -> {
+//                    for (annotation in extension.directlyApplicableAnnotations) {
+//                        extensionsWithAnnotatedMode.put(annotation, extension)
+//                    }
+//                    for (annotation in extension.childrenApplicableAnnotations) {
+//                        extensionsWithAllInAnnotatedMode.put(annotation, extension)
+//                    }
+//                }
+//                FirExtension.Mode.ALL -> extensionsWithAllMode += extension
+//            }
+//            extensionsByPlugin.put(extension.key, extension)
+//        }
+//        registerComponent(
+//            extensionClass,
+//            FirRegisteredExtension(
+//                extensionsWithAllMode,
+//                extensionsWithAnnotatedMode,
+//                extensionsWithAllInAnnotatedMode,
+//                extensionsByPlugin
+//            )
+//        )
    }

    @Deprecated("")
    fun registerUserDefinedAnnotation(metaAnnotation: AnnotationFqn, annotations: Collection<FirRegularClass>) {
-        for (annotation in annotations) {
-            require(annotation.classKind == ClassKind.ANNOTATION_CLASS)
-            val fqName = annotation.symbol.classId.asSingleFqName()
-            _annotations += fqName
-            userDefinedAnnotations.put(metaAnnotation, fqName)
-            val extensions = extensionsWithMetaAnnotations[metaAnnotation]
-            if (extensions.isEmpty()) return
-            for (extension in extensions) {
-                val registeredExtensions = this[extension.extensionType]
-
-                val metaAnnotationMode = extension.metaAnnotations.getValue(metaAnnotation)
-
-                if (metaAnnotationMode.directed) {
-                    @Suppress("UNCHECKED_CAST")
-                    (registeredExtensions.extensionsWithAnnotatedMode as Multimap<AnnotationFqn, FirExtension>).put(fqName, extension)
-                }
-
-                if (metaAnnotationMode.children) {
-                    @Suppress("UNCHECKED_CAST")
-                    (registeredExtensions.extensionsWithAllInAnnotatedMode as Multimap<AnnotationFqn, FirExtension>).put(fqName, extension)
-                }
-            }
-        }
+//        for (annotation in annotations) {
+//            require(annotation.classKind == ClassKind.ANNOTATION_CLASS)
+//            val fqName = annotation.symbol.classId.asSingleFqName()
+//            _annotations += fqName
+//            userDefinedAnnotations.put(metaAnnotation, fqName)
+//            val extensions = extensionsWithMetaAnnotations[metaAnnotation]
+//            if (extensions.isEmpty()) return
+//            for (extension in extensions) {
+//                val registeredExtensions = this[extension.extensionType]
+//
+//                val metaAnnotationMode = extension.metaAnnotations.getValue(metaAnnotation)
+//
+//                if (metaAnnotationMode.directed) {
+//                    @Suppress("UNCHECKED_CAST")
+//                    (registeredExtensions.extensionsWithAnnotatedMode as Multimap<AnnotationFqn, FirExtension>).put(fqName, extension)
+//                }
+//
+//                if (metaAnnotationMode.children) {
+//                    @Suppress("UNCHECKED_CAST")
+//                    (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
@@ -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 {
-            return FirRegisteredPluginAnnotationsImpl()
+        fun create(session: FirSession): FirRegisteredPluginAnnotations {
+            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()
@@ -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,
@@ -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.FirQualifiedAccessExpression
-import org.jetbrains.kotlin.fir.expressions.FirResolvable
-import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
+import org.jetbrains.kotlin.fir.expressions.*
 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()
+}
@@ -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.FirExtensionRegistrar
-import org.jetbrains.kotlin.fir.extensions.oldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.FirExtensionService
+import org.jetbrains.kotlin.fir.extensions.extensionService
 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) {
-        service.registerExtensions(FirExtensionRegistrar.RegisteredExtensions.EMPTY)
+    open fun registerFirExtensions(service: FirExtensionService) {
+        service.registerExtensions(emptyList())
    }

    private fun mapKtFilesToFirFiles(session: FirSession, ktFiles: List<KtFile>, firFiles: MutableList<FirFile>, useLightTree: Boolean) {
@@ -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.oldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.extensionService
 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())
    }
 }

@@ -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.oldExtensionsService
+import org.jetbrains.kotlin.fir.extensions.extensionService
 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())
            }
        }
    }
@@ -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)
        }
    }