[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/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())
     }
 }