[FIR] Correctly initialize extensions in all sessions

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

@@ -57,6 +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.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
@@ -329,7 +330,7 @@ object KotlinToJVMBytecodeCompiler {
                     dependenciesInfo, provider, librariesScope,
                     project, environment.createPackagePartProvider(librariesScope)
                 )
-                it.extensionService.registerExtensions(emptyList())
+                it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
             }
             val firProvider = (session.firProvider as FirProviderImpl)
             val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)

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

@@ -13,7 +13,13 @@ abstract class FirExtensionRegistrar {
     companion object : ProjectExtensionDescriptor<FirExtensionRegistrar>(
         "org.jetbrains.kotlin.fir.frontendIrExtension",
         FirExtensionRegistrar::class.java
-    )
+    ) {
+        val AVAILABLE_EXTENSIONS = listOf(
+            FirStatusTransformerExtension::class,
+            FirClassGenerationExtension::class,
+            AbstractFirAdditionalCheckersExtension::class
+        )
+    }
 
     protected abstract fun ExtensionRegistrarContext.configurePlugin()
 
@@ -37,26 +43,41 @@ abstract class FirExtensionRegistrar {
         }
     }
 
-    fun configure(): Collection<RegisteredExtensionsFactories<*>> {
+    @OptIn(PluginServicesInitialization::class)
+    fun configure(): BunchOfRegisteredExtensions {
         ExtensionRegistrarContext().configurePlugin()
-        return map.values
+        return BunchOfRegisteredExtensions(map.values)
     }
 
-    class RegisteredExtensionsFactories<P : FirExtension>(val kClass: KClass<P>) {
-        val extensionFactories: MutableList<FirExtension.Factory<P>> = mutableListOf()
+    class RegisteredExtensionsFactories(val kClass: KClass<out FirExtension>) {
+        val extensionFactories: MutableList<FirExtension.Factory<FirExtension>> = mutableListOf()
     }
 
-    private val map: MutableMap<KClass<out FirExtension>, RegisteredExtensionsFactories<*>> = mutableMapOf()
+    private val map: Map<KClass<out FirExtension>, RegisteredExtensionsFactories> = AVAILABLE_EXTENSIONS.map {
+        it to RegisteredExtensionsFactories(it)
+    }.toMap()
 
     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>
+        val registeredExtensions = map.getValue(kClass)
         registeredExtensions.extensionFactories += factory
     }
 }
 
+class BunchOfRegisteredExtensions @PluginServicesInitialization constructor(
+    val extensions: Collection<FirExtensionRegistrar.RegisteredExtensionsFactories>
+) {
+    companion object {
+        @OptIn(PluginServicesInitialization::class)
+        fun empty(): BunchOfRegisteredExtensions {
+            val extensions = FirExtensionRegistrar.AVAILABLE_EXTENSIONS.map { FirExtensionRegistrar.RegisteredExtensionsFactories(it) }
+            return BunchOfRegisteredExtensions(extensions)
+        }
+    }
+}
+
 @OptIn(PluginServicesInitialization::class)
-fun FirExtensionService.registerExtensions(registeredExtensions: Collection<FirExtensionRegistrar.RegisteredExtensionsFactories<*>>) {
-    registeredExtensions.forEach { registerExtensions(it.kClass, it.extensionFactories) }
+fun FirExtensionService.registerExtensions(registeredExtensions: BunchOfRegisteredExtensions) {
+    registeredExtensions.extensions.forEach { registerExtensions(it.kClass, it.extensionFactories) }
     session.registeredPluginAnnotations.initialize()
 }

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

@@ -26,6 +26,7 @@ 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.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
@@ -109,7 +110,7 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
     }
 
     open fun registerFirExtensions(service: FirExtensionService) {
-        service.registerExtensions(emptyList())
+        service.registerExtensions(BunchOfRegisteredExtensions.empty())
     }
 
     private fun mapKtFilesToFirFiles(session: FirSession, ktFiles: List<KtFile>, firFiles: MutableList<FirFile>, useLightTree: Boolean) {

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

@@ -8,6 +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.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
@@ -34,7 +35,7 @@ fun createSession(
     val provider = FirProjectSessionProvider(project)
     return FirJavaModuleBasedSession(moduleInfo, provider, sourceScope).also {
         createSessionForDependencies(project, provider, moduleInfo, librariesScope, packagePartProvider)
-        it.extensionService.registerExtensions(emptyList())
+        it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
     }
 }
 

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

@@ -14,6 +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.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
@@ -39,7 +40,7 @@ interface FirModuleResolveState {
                 sessionProvider.sessionCache[moduleInfo] = moduleBasedSession
             }
             session.also {
-                it.extensionService.registerExtensions(emptyList())
+                it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
             }
         }
     }

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.dependenciesWithoutSelf
+import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
 import org.jetbrains.kotlin.fir.extensions.extensionService
 import org.jetbrains.kotlin.fir.extensions.registerExtensions
 import org.jetbrains.kotlin.fir.java.FirJavaModuleBasedSession
@@ -74,7 +75,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.extensionService.registerExtensions(emptyList())
+            it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
         }
     }