[FIR] Add ability to register session components from FIR plugins

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

@@ -30,6 +30,7 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             FirSupertypeGenerationExtension::class,
             FirTypeAttributeExtension::class,
             FirExpressionResolutionExtension::class,
+            FirExtensionSessionComponent::class,
         )
     }
 
@@ -68,6 +69,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             registerExtension(FirExpressionResolutionExtension::class, this)
         }
 
+        @JvmName("plusExtensionSessionComponent")
+        operator fun (FirExtensionSessionComponent.Factory).unaryPlus() {
+            registerExtension(FirExtensionSessionComponent::class, this)
+        }
+
         // ------------------ reference methods ------------------
 
         @JvmName("plusStatusTransformerExtension")
@@ -99,6 +105,11 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
         operator fun ((FirSession) -> FirExpressionResolutionExtension).unaryPlus() {
             FirExpressionResolutionExtension.Factory { this.invoke(it) }.unaryPlus()
         }
+
+        @JvmName("plusExtensionSessionComponent")
+        operator fun ((FirSession) -> FirExtensionSessionComponent).unaryPlus() {
+            FirExtensionSessionComponent.Factory { this.invoke(it) }.unaryPlus()
+        }
     }
 
     @OptIn(PluginServicesInitialization::class)
@@ -125,7 +136,6 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
     private var isInitialized: AtomicBoolean = AtomicBoolean(false)
 
     private fun <P : FirExtension> registerExtension(kClass: KClass<out P>, factory: FirExtension.Factory<P>) {
-        @Suppress("UNCHECKED_CAST")
         val registeredExtensions = map.getValue(kClass)
         registeredExtensions.extensionFactories += factory
     }
@@ -152,5 +162,8 @@ class BunchOfRegisteredExtensions @PluginServicesInitialization constructor(
 @OptIn(PluginServicesInitialization::class)
 fun FirExtensionService.registerExtensions(registeredExtensions: BunchOfRegisteredExtensions) {
     registeredExtensions.extensions.forEach { registerExtensions(it.kClass, it.extensionFactories) }
+    extensionSessionComponents.forEach {
+        session.register(it.componentClass, it)
+    }
     session.registeredPluginAnnotations.initialize()
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/extensions/FirExtensionSessionComponent.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2022 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 kotlin.reflect.KClass
+
+abstract class FirExtensionSessionComponent(session: FirSession) : FirExtension(session), FirSessionComponent {
+    companion object {
+        val NAME = FirExtensionPointName("ExtensionSessionComponent")
+    }
+
+    final override val name: FirExtensionPointName
+        get() = NAME
+
+    final override val extensionType: KClass<out FirExtension>
+        get() = FirExtensionSessionComponent::class
+
+    open val componentClass: KClass<out FirExtensionSessionComponent>
+        get() = this::class
+
+    fun interface Factory : FirExtension.Factory<FirExtensionSessionComponent>
+}
+
+val FirExtensionService.extensionSessionComponents: List<FirExtensionSessionComponent> by FirExtensionService.registeredExtensions()