[FIR] Fix registration of multiple compiler plugins

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

@@ -142,46 +142,48 @@ abstract class FirExtensionRegistrar : FirExtensionRegistrarAdapter() {
             ExtensionRegistrarContext().configurePlugin()
         }
 
-        return BunchOfRegisteredExtensions(map.values)
+        return BunchOfRegisteredExtensions(map)
     }
 
-    class RegisteredExtensionsFactories(val kClass: KClass<out FirExtension>) {
+    private val map: Map<KClass<out FirExtension>, MutableList<FirExtension.Factory<FirExtension>>> = AVAILABLE_EXTENSIONS.associateWith {
-        val extensionFactories: MutableList<FirExtension.Factory<FirExtension>> = mutableListOf()
+        mutableListOf()
-    }
-
-    private val map: Map<KClass<out FirExtension>, RegisteredExtensionsFactories> = AVAILABLE_EXTENSIONS.associateWith {
-        RegisteredExtensionsFactories(it)
     }
 
     private var isInitialized: AtomicBoolean = AtomicBoolean(false)
 
     private fun <P : FirExtension> registerExtension(kClass: KClass<out P>, factory: FirExtension.Factory<P>) {
         val registeredExtensions = map.getValue(kClass)
-        registeredExtensions.extensionFactories += factory
+        registeredExtensions += factory
     }
 }
 
 class BunchOfRegisteredExtensions @PluginServicesInitialization constructor(
-    val extensions: Collection<FirExtensionRegistrar.RegisteredExtensionsFactories>
+    val extensions: Map<KClass<out FirExtension>, List<FirExtension.Factory<FirExtension>>>
 ) {
     companion object {
         @OptIn(PluginServicesInitialization::class)
         fun empty(): BunchOfRegisteredExtensions {
-            val extensions = FirExtensionRegistrar.AVAILABLE_EXTENSIONS.map { FirExtensionRegistrar.RegisteredExtensionsFactories(it) }
+            return BunchOfRegisteredExtensions(FirExtensionRegistrar.AVAILABLE_EXTENSIONS.associateWith { listOf() })
-            return BunchOfRegisteredExtensions(extensions)
         }
     }
 
     @OptIn(PluginServicesInitialization::class)
     operator fun plus(other: BunchOfRegisteredExtensions): BunchOfRegisteredExtensions {
-        return BunchOfRegisteredExtensions(extensions + other.extensions)
+        val combinedExtensions = buildMap {
+            for (extensionClass in FirExtensionRegistrar.AVAILABLE_EXTENSIONS) {
+                put(extensionClass, extensions.getValue(extensionClass) + other.extensions.getValue(extensionClass))
+            }
+        }
+        return BunchOfRegisteredExtensions(combinedExtensions)
     }
 }
 
 @SessionConfiguration
 @OptIn(PluginServicesInitialization::class)
 fun FirExtensionService.registerExtensions(registeredExtensions: BunchOfRegisteredExtensions) {
-    registeredExtensions.extensions.forEach { registerExtensions(it.kClass, it.extensionFactories) }
+    registeredExtensions.extensions.forEach { (extensionClass, extensionFactories) ->
+        registerExtensions(extensionClass, extensionFactories)
+    }
     extensionSessionComponents.forEach {
         session.register(it.componentClass, it)
     }

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

@@ -105,9 +105,6 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
                 needRegisterJavaElementFinder = true
             ) {
                 configureSession()
-                getFirExtensions()?.let {
-                    registerExtensions(it)
-                }
             }
         }
 
@@ -127,10 +124,6 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         runAnalysis(testDataFile, files, firFilesPerSession)
     }
 
-    protected open fun getFirExtensions(): BunchOfRegisteredExtensions? {
-        return null
-    }
-
     private fun mapKtFilesToFirFiles(session: FirSession, ktFiles: List<KtFile>, firFiles: MutableList<FirFile>, useLightTree: Boolean) {
         val firProvider = (session.firProvider as FirProviderImpl)
         if (useLightTree) {