[FIR] Refactor session components initialization

From now all customization of session should perform via DSL which
  declared in FirSessionFactory. Registration or mutation session
  services from other places is prohibited

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

@@ -64,15 +64,12 @@ import org.jetbrains.kotlin.fir.backend.jvm.FirJvmClassCodegen
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmKotlinMangler
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmVisibilityConverter
 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.checkers.registerExtendedCommonCheckers
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.resolve.firProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirProviderImpl
 import org.jetbrains.kotlin.fir.resolve.transformers.FirTotalResolveProcessor
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
-import org.jetbrains.kotlin.fir.session.registerExtendedCheckersComponent
 import org.jetbrains.kotlin.idea.MainFunctionDetector
 import org.jetbrains.kotlin.ir.backend.jvm.jvmResolveLibraries
 import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmManglerDesc
@@ -353,7 +350,11 @@ object KotlinToJVMBytecodeCompiler {
             }
 
             val moduleInfo = FirJvmModuleInfo(module.getModuleName())
-            val session: FirSession = FirSessionFactory.createJavaModuleBasedSession(moduleInfo, provider, scope).also {
+            val session: FirSession = FirSessionFactory.createJavaModuleBasedSession(moduleInfo, provider, scope) {
+                if (extendedAnalysisMode) {
+                    registerExtendedCommonCheckers()
+                }
+            }.also {
                 val dependenciesInfo = FirJvmModuleInfo(Name.special("<dependencies>"))
                 moduleInfo.dependencies.add(dependenciesInfo)
                 val librariesScope = ProjectScope.getLibrariesScope(project)
@@ -361,10 +362,6 @@ object KotlinToJVMBytecodeCompiler {
                     dependenciesInfo, provider, librariesScope,
                     project, environment.createPackagePartProvider(librariesScope)
                 )
-                it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
-                if (extendedAnalysisMode) {
-                    it.registerExtendedCheckersComponent()
-                }
             }
             val firProvider = (session.firProvider as FirProviderImpl)
             val builder = RawFirBuilder(session, firProvider.kotlinScopeProvider, stubMode = false)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/CheckersComponent.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.analysis
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
+import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.cfa.AbstractFirPropertyInitializationChecker
 import org.jetbrains.kotlin.fir.analysis.checkers.cfa.FirControlFlowChecker
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.*
@@ -14,29 +15,23 @@ import org.jetbrains.kotlin.fir.analysis.checkers.expression.*
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
 
 class CheckersComponent : FirSessionComponent {
-    companion object {
-        fun componentWithDefaultCheckers(): CheckersComponent {
-            return CheckersComponent().apply {
-                register(CommonDeclarationCheckers)
-                register(CommonExpressionCheckers)
-            }
-        }
-    }
-
     val declarationCheckers: DeclarationCheckers get() = _declarationCheckers
     private val _declarationCheckers = ComposedDeclarationCheckers()
 
     val expressionCheckers: ExpressionCheckers get() = _expressionCheckers
     private val _expressionCheckers = ComposedExpressionCheckers()
 
+    @SessionConfiguration
     fun register(checkers: DeclarationCheckers) {
         _declarationCheckers.register(checkers)
     }
 
+    @SessionConfiguration
     fun register(checkers: ExpressionCheckers) {
         _expressionCheckers.register(checkers)
     }
 
+    @SessionConfiguration
     fun register(checkers: FirAdditionalCheckersExtension) {
         register(checkers.declarationCheckers)
         register(checkers.expressionCheckers)

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/DefaultExpressionCheckers.kt

@@ -5,20 +5,3 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.expression
 
-import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSealedClassConstructorCallChecker
-
-object CommonExpressionCheckers : ExpressionCheckers() {
-    override val expressionCheckers: List<FirBasicExpresionChecker> = listOf()
-    override val qualifiedAccessCheckers: List<FirQualifiedAccessChecker> = listOf(
-        FirSuperNotAvailableChecker,
-        FirNotASupertypeChecker,
-        FirSuperclassNotAccessibleFromInterfaceChecker,
-        FirAbstractSuperCallChecker,
-        FirQualifiedSupertypeExtendedByOtherSupertypeChecker,
-        FirProjectionsOnNonClassTypeArgumentChecker,
-        FirUpperBoundViolatedChecker,
-        FirTypeArgumentsNotAllowedExpressionChecker,
-        FirSealedClassConstructorCallChecker,
-    )
-    override val functionCallCheckers: List<FirFunctionCallChecker> = listOf()
-}

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/collectors/FirDiagnosticsCollector.kt

@@ -6,13 +6,9 @@
 package org.jetbrains.kotlin.fir.analysis.collectors
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.analysis.checkersComponent
-import org.jetbrains.kotlin.fir.analysis.extensions.additionalCheckers
-import org.jetbrains.kotlin.fir.extensions.extensionService
 
 object FirDiagnosticsCollector {
     fun create(session: FirSession): AbstractDiagnosticCollector {
-        session.registerAdditionalCheckers()
         val collector = SimpleDiagnosticsCollector(session)
         collector.registerAllComponents()
         return collector
@@ -21,13 +17,8 @@ object FirDiagnosticsCollector {
     // Use in CLI compiler
     @Suppress("unused")
     fun createParallel(session: FirSession): AbstractDiagnosticCollector {
-        session.registerAdditionalCheckers()
         val collector = ParallelDiagnosticsCollector(session, numberOfThreads = 4)
         collector.registerAllComponents()
         return collector
     }
-
-    private fun FirSession.registerAdditionalCheckers() {
-        extensionService.additionalCheckers.forEach(checkersComponent::register)
-    }
-}
+}

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/checkers/CheckersContainers.kt

@@ -0,0 +1,20 @@
+/*
+ * 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.checkers
+
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.ExtendedDeclarationCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExtendedExpressionCheckers
+import org.jetbrains.kotlin.fir.session.FirSessionFactory
+
+fun FirSessionFactory.FirSessionConfigurator.registerCommonCheckers() {
+    useCheckers(CommonDeclarationCheckers)
+    useCheckers(CommonExpressionCheckers)
+}
+
+fun FirSessionFactory.FirSessionConfigurator.registerExtendedCommonCheckers() {
+    useCheckers(ExtendedExpressionCheckers)
+    useCheckers(ExtendedDeclarationCheckers)
+}

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/checkers/CommonDeclarationCheckers.kt

@@ -3,18 +3,17 @@
  * 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.analysis.checkers.declaration
+package org.jetbrains.kotlin.fir.checkers
 
 import org.jetbrains.kotlin.fir.analysis.cfa.AbstractFirPropertyInitializationChecker
 import org.jetbrains.kotlin.fir.analysis.cfa.FirCallsEffectAnalyzer
 import org.jetbrains.kotlin.fir.analysis.cfa.FirPropertyInitializationAnalyzer
 import org.jetbrains.kotlin.fir.analysis.cfa.FirReturnsImpliesAnalyzer
 import org.jetbrains.kotlin.fir.analysis.checkers.cfa.FirControlFlowChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.*
 
 object CommonDeclarationCheckers : DeclarationCheckers() {
-    override val fileCheckers: List<FirFileChecker> = listOf(
-
-    )
+    override val fileCheckers: List<FirFileChecker> = listOf()
 
     override val declarationCheckers: List<FirBasicDeclarationChecker> = listOf(
         FirAnnotationClassDeclarationChecker,
@@ -27,7 +26,7 @@ object CommonDeclarationCheckers : DeclarationCheckers() {
         FirConflictingProjectionChecker,
     )
 
-    override val memberDeclarationCheckers: List<FirMemberDeclarationChecker> =listOf(
+    override val memberDeclarationCheckers: List<FirMemberDeclarationChecker> = listOf(
         FirInfixFunctionDeclarationChecker,
         FirExposedVisibilityDeclarationChecker,
         FirCommonConstructorDelegationIssuesChecker,
@@ -59,4 +58,4 @@ object CommonDeclarationCheckers : DeclarationCheckers() {
     override val variableAssignmentCfaBasedCheckers: List<AbstractFirPropertyInitializationChecker> = listOf(
         FirPropertyInitializationAnalyzer
     )
-}
+}

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/checkers/CommonExpressionCheckers.kt

@@ -0,0 +1,25 @@
+/*
+ * 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.checkers
+
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirSealedClassConstructorCallChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.*
+
+object CommonExpressionCheckers : ExpressionCheckers() {
+    override val expressionCheckers: List<FirBasicExpresionChecker> = listOf()
+    override val qualifiedAccessCheckers: List<FirQualifiedAccessChecker> = listOf(
+        FirSuperNotAvailableChecker,
+        FirNotASupertypeChecker,
+        FirSuperclassNotAccessibleFromInterfaceChecker,
+        FirAbstractSuperCallChecker,
+        FirQualifiedSupertypeExtendedByOtherSupertypeChecker,
+        FirProjectionsOnNonClassTypeArgumentChecker,
+        FirUpperBoundViolatedChecker,
+        FirTypeArgumentsNotAllowedExpressionChecker,
+        FirSealedClassConstructorCallChecker,
+    )
+    override val functionCallCheckers: List<FirFunctionCallChecker> = listOf()
+}

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt

@@ -5,30 +5,25 @@
 
 package org.jetbrains.kotlin.fir.session
 
-import org.jetbrains.kotlin.fir.FirModuleBasedSession
-import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.CheckersComponent
-import org.jetbrains.kotlin.fir.analysis.checkers.declaration.ExtendedDeclarationCheckers
-import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExtendedExpressionCheckers
-import org.jetbrains.kotlin.fir.analysis.checkersComponent
 import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
-import org.jetbrains.kotlin.fir.java.registerJavaVisibilityChecker
-import org.jetbrains.kotlin.fir.registerJvmEffectiveVisibilityResolver
-import org.jetbrains.kotlin.fir.resolve.FirQualifierResolver
-import org.jetbrains.kotlin.fir.resolve.FirTypeResolver
-import org.jetbrains.kotlin.fir.resolve.calls.jvm.registerJvmCallConflictResolverFactory
+import org.jetbrains.kotlin.fir.java.FirJavaVisibilityChecker
+import org.jetbrains.kotlin.fir.resolve.*
+import org.jetbrains.kotlin.fir.resolve.calls.ConeCallConflictResolverFactory
+import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticNamesProvider
+import org.jetbrains.kotlin.fir.resolve.calls.jvm.JvmCallConflictResolverFactory
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirQualifierResolverImpl
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirTypeResolverImpl
-import org.jetbrains.kotlin.fir.resolve.registerJavaClassMapper
-import org.jetbrains.kotlin.fir.resolve.registerJavaSyntheticNamesProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.plugin.GeneratedClassIndex
 import org.jetbrains.kotlin.fir.scopes.impl.FirDeclaredMemberScopeProvider
 import org.jetbrains.kotlin.fir.types.FirCorrespondingSupertypesCache
 
 // -------------------------- Required components --------------------------
 
+@OptIn(SessionConfiguration::class)
 fun FirSession.registerCommonComponents() {
     register(FirDeclaredMemberScopeProvider::class, FirDeclaredMemberScopeProvider())
     register(FirCorrespondingSupertypesCache::class, FirCorrespondingSupertypesCache(this))
@@ -41,31 +36,24 @@ fun FirSession.registerCommonComponents() {
 
 // -------------------------- Resolve components --------------------------
 
-// TODO: rename to `registerCommonResolveComponents
+/*
+ * Resolve components which are same on all platforms
+ */
+@OptIn(SessionConfiguration::class)
 fun FirModuleBasedSession.registerResolveComponents() {
     register(FirQualifierResolver::class, FirQualifierResolverImpl(this))
     register(FirTypeResolver::class, FirTypeResolverImpl(this))
+    register(CheckersComponent::class, CheckersComponent())
 }
 
-fun FirModuleBasedSession.registerJavaSpecificComponents() {
-    registerJavaVisibilityChecker()
-    registerJvmCallConflictResolverFactory()
-    registerJvmEffectiveVisibilityResolver()
-    registerJavaClassMapper()
-    registerJavaSyntheticNamesProvider()
-}
-
-// -------------------------- Checker components --------------------------
-
 /*
- * TODO: in future rename to `registerCheckersComponent` and configure
- *    exact checkers according to platforms of current session
+ * Resolve components which have specific implementations on JVM
  */
-fun FirModuleBasedSession.registerCheckersComponent() {
-    register(CheckersComponent::class, CheckersComponent.componentWithDefaultCheckers())
-}
-
-fun FirModuleBasedSession.registerExtendedCheckersComponent() {
-    this.checkersComponent.register(ExtendedExpressionCheckers)
-    this.checkersComponent.register(ExtendedDeclarationCheckers)
+@OptIn(SessionConfiguration::class)
+fun FirModuleBasedSession.registerJavaSpecificComponents() {
+    register(FirVisibilityChecker::class, FirJavaVisibilityChecker)
+    register(ConeCallConflictResolverFactory::class, JvmCallConflictResolverFactory)
+    register(FirEffectiveVisibilityResolver::class, FirJvmEffectiveVisibilityResolver(this))
+    register(FirPlatformClassMapper::class, FirJavaClassMapper(this))
+    register(FirSyntheticNamesProvider::class, FirJavaSyntheticNamesProvider)
 }

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirSessionFactory.kt

@@ -10,8 +10,18 @@ import com.intellij.psi.PsiElementFinder
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.fir.FirModuleBasedSession
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
+import org.jetbrains.kotlin.fir.SessionConfiguration
+import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
+import org.jetbrains.kotlin.fir.analysis.checkersComponent
+import org.jetbrains.kotlin.fir.analysis.extensions.additionalCheckers
+import org.jetbrains.kotlin.fir.checkers.registerCommonCheckers
+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.*
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
@@ -23,18 +33,40 @@ import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory
 
-@OptIn(PrivateSessionConstructor::class)
+@OptIn(PrivateSessionConstructor::class, SessionConfiguration::class)
 object FirSessionFactory {
+    class FirSessionConfigurator(private val session: FirModuleBasedSession) {
+        private val registeredExtensions = mutableListOf<BunchOfRegisteredExtensions>(BunchOfRegisteredExtensions.empty())
+
+        fun registerExtensions(extensions: BunchOfRegisteredExtensions) {
+            registeredExtensions += extensions
+        }
+
+        fun useCheckers(checkers: ExpressionCheckers) {
+            session.checkersComponent.register(checkers)
+        }
+
+        fun useCheckers(checkers: DeclarationCheckers) {
+            session.checkersComponent.register(checkers)
+        }
+
+        @SessionConfiguration
+        fun configure() {
+            session.extensionService.registerExtensions(registeredExtensions.reduce(BunchOfRegisteredExtensions::plus))
+            session.extensionService.additionalCheckers.forEach(session.checkersComponent::register)
+        }
+    }
+
     fun createJavaModuleBasedSession(
         moduleInfo: ModuleInfo,
         sessionProvider: FirProjectSessionProvider,
         scope: GlobalSearchScope,
-        dependenciesProvider: FirSymbolProvider? = null
+        dependenciesProvider: FirSymbolProvider? = null,
+        init: FirSessionConfigurator.() -> Unit = {}
     ): FirJavaModuleBasedSession {
         return FirJavaModuleBasedSession(moduleInfo, sessionProvider).apply {
             registerCommonComponents()
             registerResolveComponents()
-            registerCheckersComponent()
             registerJavaSpecificComponents()
 
             val kotlinScopeProvider = KotlinScopeProvider(::wrapScopeWithJvmMapped)
@@ -54,6 +86,11 @@ object FirSessionFactory {
                 ) as FirSymbolProvider
             )
 
+            FirSessionConfigurator(this).apply {
+                registerCommonCheckers()
+                init()
+            }.configure()
+
             PsiElementFinder.EP.getPoint(sessionProvider.project)
                 .registerExtension(FirJavaElementFinder(this, sessionProvider.project), sessionProvider.project)
         }

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

@@ -42,7 +42,3 @@ object FirJavaVisibilityChecker : FirVisibilityChecker() {
         }
     }
 }
-
-fun FirSession.registerJavaVisibilityChecker() {
-    register(FirVisibilityChecker::class, FirJavaVisibilityChecker)
-}

compiler/fir/jvm/src/org/jetbrains/kotlin/fir/FirJvmEffectiveVisibilityResolver.kt

@@ -15,7 +15,3 @@ class FirJvmEffectiveVisibilityResolver(session: FirSession) : FirEffectiveVisib
         }
     }
 }
-
-fun FirSession.registerJvmEffectiveVisibilityResolver() {
-    register(FirEffectiveVisibilityResolver::class, FirJvmEffectiveVisibilityResolver(this))
-}

compiler/fir/jvm/src/org/jetbrains/kotlin/fir/resolve/FirJavaClassMapper.kt

@@ -16,7 +16,3 @@ class FirJavaClassMapper(private val session: FirSession) : FirPlatformClassMapp
         return javaClassId?.let { session.firSymbolProvider.getClassLikeSymbolByFqName(it)?.fir } as? FirRegularClass
     }
 }
-
-fun FirSession.registerJavaClassMapper() {
-    register(FirPlatformClassMapper::class, FirJavaClassMapper(this))
-}

compiler/fir/jvm/src/org/jetbrains/kotlin/fir/resolve/FirJavaSyntheticNamesProvider.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve
 
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticNamesProvider
 import org.jetbrains.kotlin.load.java.propertyNameByGetMethodName
 import org.jetbrains.kotlin.name.Name
@@ -41,6 +40,3 @@ object FirJavaSyntheticNamesProvider : FirSyntheticNamesProvider() {
     }
 }
 
-fun FirSession.registerJavaSyntheticNamesProvider() {
-    register(FirSyntheticNamesProvider::class, FirJavaSyntheticNamesProvider)
-}

compiler/fir/jvm/src/org/jetbrains/kotlin/fir/resolve/calls/jvm/JvmCallConflictResolverFactory.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.calls.jvm
 
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.calls.ConeCallConflictResolverFactory
 import org.jetbrains.kotlin.fir.resolve.calls.ConeCompositeConflictResolver
 import org.jetbrains.kotlin.fir.resolve.calls.ConeOverloadConflictResolver
@@ -25,7 +24,3 @@ object JvmCallConflictResolverFactory : ConeCallConflictResolverFactory() {
         )
     }
 }
-
-fun FirSession.registerJvmCallConflictResolverFactory() {
-    register(ConeCallConflictResolverFactory::class, JvmCallConflictResolverFactory)
-}

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.extensions
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.SessionConfiguration
 import kotlin.reflect.KClass
 
 abstract class FirExtensionRegistrar {
@@ -76,8 +77,14 @@ class BunchOfRegisteredExtensions @PluginServicesInitialization constructor(
             return BunchOfRegisteredExtensions(extensions)
         }
     }
+
+    @OptIn(PluginServicesInitialization::class)
+    operator fun plus(other: BunchOfRegisteredExtensions): BunchOfRegisteredExtensions {
+        return BunchOfRegisteredExtensions(extensions + other.extensions)
+    }
 }
 
+@SessionConfiguration
 @OptIn(PluginServicesInitialization::class)
 fun FirExtensionService.registerExtensions(registeredExtensions: BunchOfRegisteredExtensions) {
     registeredExtensions.extensions.forEach { registerExtensions(it.kClass, it.extensionFactories) }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt

@@ -18,6 +18,9 @@ import kotlin.reflect.KClass
 @RequiresOptIn
 annotation class PrivateSessionConstructor
 
+@RequiresOptIn
+annotation class SessionConfiguration
+
 interface FirSessionComponent
 
 abstract class FirSession @PrivateSessionConstructor constructor(val sessionProvider: FirSessionProvider?) : ComponentArrayOwner<FirSessionComponent, FirSessionComponent>() {
@@ -39,6 +42,7 @@ abstract class FirSession @PrivateSessionConstructor constructor(val sessionProv
 
     final override val typeRegistry: TypeRegistry<FirSessionComponent, FirSessionComponent> = Companion
 
+    @SessionConfiguration
     fun register(tClass: KClass<out FirSessionComponent>, value: FirSessionComponent) {
         registerComponent(tClass, value)
     }

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

@@ -5,13 +5,11 @@
 
 package org.jetbrains.kotlin.fir
 
-import org.jetbrains.kotlin.fir.analysis.checkers.declaration.ExtendedDeclarationCheckers
-import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExtendedExpressionCheckers
-import org.jetbrains.kotlin.fir.analysis.checkersComponent
+import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
+import org.jetbrains.kotlin.fir.session.FirSessionFactory
 
 abstract class AbstractExtendedFirDiagnosticsTest : AbstractFirDiagnosticsTest() {
-    override fun configureSession(session: FirSession) {
-        session.checkersComponent.register(ExtendedDeclarationCheckers)
-        session.checkersComponent.register(ExtendedExpressionCheckers)
+    override fun FirSessionFactory.FirSessionConfigurator.configureSession() {
+        registerExtendedCommonCheckers()
     }
-}
+}

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

@@ -5,13 +5,11 @@
 
 package org.jetbrains.kotlin.fir
 
-import org.jetbrains.kotlin.fir.analysis.checkers.declaration.ExtendedDeclarationCheckers
-import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExtendedExpressionCheckers
-import org.jetbrains.kotlin.fir.analysis.checkersComponent
+import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
+import org.jetbrains.kotlin.fir.session.FirSessionFactory
 
 abstract class AbstractExtendedFirWithLightTreeDiagnosticsTest : AbstractFirDiagnosticsWithLightTreeTest() {
-    override fun configureSession(session: FirSession) {
-        session.checkersComponent.register(ExtendedDeclarationCheckers)
-        session.checkersComponent.register(ExtendedExpressionCheckers)
+    override fun FirSessionFactory.FirSessionConfigurator.configureSession() {
+        registerExtendedCommonCheckers()
     }
-}
+}

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

@@ -27,9 +27,6 @@ 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
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.lightTree.LightTree2Fir
 import org.jetbrains.kotlin.fir.resolve.firProvider
@@ -79,18 +76,18 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         FirSessionFactory.createLibrarySession(
             builtInsModuleInfo, sessionProvider, allProjectScope, project,
             environment.createPackagePartProvider(allProjectScope)
-        ).also {
-            registerFirExtensions(it.extensionService)
-        }
+        )
 
         val configToSession = modules.mapValues { (config, info) ->
             val moduleFiles = groupedByModule.getValue(config)
             val scope = TopDownAnalyzerFacadeForJVM.newModuleSearchScope(
                 project,
                 moduleFiles.mapNotNull { it.ktFile })
-            FirSessionFactory.createJavaModuleBasedSession(info, sessionProvider, scope).also {
-                registerFirExtensions(it.extensionService)
-                configureSession(it)
+            FirSessionFactory.createJavaModuleBasedSession(info, sessionProvider, scope) {
+                configureSession()
+                getFirExtensions()?.let {
+                    registerExtensions(it)
+                }
             }
         }
 
@@ -110,8 +107,8 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         runAnalysis(testDataFile, files, firFilesPerSession)
     }
 
-    open fun registerFirExtensions(service: FirExtensionService) {
-        service.registerExtensions(BunchOfRegisteredExtensions.empty())
+    protected open fun getFirExtensions(): BunchOfRegisteredExtensions? {
+        return null
     }
 
     private fun mapKtFilesToFirFiles(session: FirSession, ktFiles: List<KtFile>, firFiles: MutableList<FirFile>, useLightTree: Boolean) {
@@ -333,5 +330,5 @@ abstract class AbstractFirBaseDiagnosticsTest : BaseDiagnosticsTest() {
         return result
     }
 
-    protected open fun configureSession(session: FirSession) {}
+    protected open fun FirSessionFactory.FirSessionConfigurator.configureSession() {}
 }

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

@@ -8,9 +8,6 @@ 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.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.session.FirSessionFactory
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
@@ -34,7 +31,6 @@ fun createSession(
     val provider = FirProjectSessionProvider(project)
     return FirSessionFactory.createJavaModuleBasedSession(moduleInfo, provider, sourceScope).also {
         createSessionForDependencies(project, provider, moduleInfo, librariesScope, packagePartProvider)
-        it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
     }
 }
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeLibrarySession.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.sessions
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
+import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
@@ -34,6 +35,7 @@ internal class FirIdeLibrariesSession private constructor(
     override val scope: GlobalSearchScope
 ) : FirIdeSession(moduleInfo, sessionProvider) {
     companion object {
+        @OptIn(SessionConfiguration::class)
         fun create(
             moduleInfo: ModuleSourceInfo,
             sessionProvider: FirIdeSessionProvider,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSourcesSession.kt

@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCompositeSymbolProvide
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.resolve.transformers.PhasedFirFileResolver
 import org.jetbrains.kotlin.fir.scopes.KotlinScopeProvider
-import org.jetbrains.kotlin.fir.session.registerCheckersComponent
 import org.jetbrains.kotlin.fir.session.registerCommonComponents
 import org.jetbrains.kotlin.fir.session.registerJavaSpecificComponents
 import org.jetbrains.kotlin.fir.session.registerResolveComponents
@@ -51,6 +50,7 @@ internal class FirIdeSourcesSession private constructor(
         /**
          * Should be invoked only under a [moduleInfo]-based lock
          */
+        @OptIn(SessionConfiguration::class)
         fun create(
             project: Project,
             moduleInfo: ModuleSourceInfo,
@@ -68,7 +68,6 @@ internal class FirIdeSourcesSession private constructor(
 
                 registerCommonComponents()
                 registerResolveComponents()
-                registerCheckersComponent()
 
                 val provider = FirIdeProvider(
                     project,

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirMultiModuleResolveTest.kt

@@ -15,9 +15,6 @@ import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 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.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.firProvider
@@ -60,9 +57,7 @@ abstract class AbstractFirMultiModuleResolveTest : AbstractMultiModuleTest() {
 
     private fun createSession(module: Module, provider: FirProjectSessionProvider): FirJavaModuleBasedSession {
         val moduleInfo = module.productionSourceInfo()!!
-        return FirSessionFactory.createJavaModuleBasedSession(moduleInfo, provider, moduleInfo.contentScope()).also {
-            it.extensionService.registerExtensions(BunchOfRegisteredExtensions.empty())
-        }
+        return FirSessionFactory.createJavaModuleBasedSession(moduleInfo, provider, moduleInfo.contentScope())
     }
 
     private fun createLibrarySession(moduleInfo: IdeaModuleInfo, provider: FirProjectSessionProvider): FirLibrarySession {

plugins/fir/fir-plugin-prototype/tests/org/jetbrains/kotlin/fir/plugin/AbstractFirAllOpenDiagnosticTest.kt

@@ -8,16 +8,15 @@ package org.jetbrains.kotlin.fir.plugin
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoot
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.fir.AbstractFirDiagnosticsTest
-import org.jetbrains.kotlin.fir.extensions.FirExtensionService
-import org.jetbrains.kotlin.fir.extensions.registerExtensions
+import org.jetbrains.kotlin.fir.extensions.BunchOfRegisteredExtensions
 import java.io.File
 
 abstract class AbstractFirAllOpenDiagnosticTest : AbstractFirDiagnosticsTest() {
     override val pluginPhasesEnabled: Boolean
         get() = true
 
-    override fun registerFirExtensions(service: FirExtensionService) {
-        service.registerExtensions(FirAllOpenComponentRegistrar().configure())
+    override fun getFirExtensions(): BunchOfRegisteredExtensions {
+        return FirAllOpenComponentRegistrar().configure()
     }
 
     override fun updateConfiguration(configuration: CompilerConfiguration) {