[LL API] Replace configuration lambdas with extensions

Caching 'LLFirSession' that is modified in arbitrary way is not safe.
The new approach is less ad-hoc, as extensions are applied to all
sessions.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/LLFirResolveStateService.kt

@@ -10,7 +10,6 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootModificationTracker
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionProviderStorage
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirLibraryOrLibrarySourceResolvableResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirNotUnderContentRootResolveSession
@@ -46,10 +45,9 @@ internal class LLFirResolveSessionService(project: Project) {
 
         internal fun createFirResolveSessionFor(
             useSiteKtModule: KtModule,
-            sessionProviderStorage: LLFirSessionProviderStorage,
-            configureSession: (LLFirSession.() -> Unit)? = null,
+            sessionProviderStorage: LLFirSessionProviderStorage
         ): LLFirResolvableResolveSession {
-            val sessionProvider = sessionProviderStorage.getSessionProvider(useSiteKtModule, configureSession)
+            val sessionProvider = sessionProviderStorage.getSessionProvider(useSiteKtModule)
             val useSiteSession = sessionProvider.rootModuleSession
             return when (useSiteKtModule) {
                 is KtSourceModule -> {
@@ -100,15 +98,11 @@ internal class LLFirResolveSessionService(project: Project) {
 }
 
 @TestOnly
-fun createFirResolveSessionForNoCaching(
-    useSiteKtModule: KtModule,
-    project: Project,
-    configureSession: (LLFirSession.() -> Unit)? = null,
-): LLFirResolveSession =
-    LLFirResolveSessionService.createFirResolveSessionFor(
+fun createFirResolveSessionForNoCaching(useSiteKtModule: KtModule, project: Project): LLFirResolveSession {
+    return LLFirResolveSessionService.createFirResolveSessionFor(
         useSiteKtModule = useSiteKtModule,
-        sessionProviderStorage = LLFirSessionProviderStorage(project),
-        configureSession = configureSession
+        sessionProviderStorage = LLFirSessionProviderStorage(project)
     )
+}
 
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionConfigurator.kt

@@ -0,0 +1,17 @@
+/*
+ * Copyright 2010-2023 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.analysis.low.level.api.fir.sessions
+
+import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
+
+interface LLFirSessionConfigurator {
+    companion object : ProjectExtensionDescriptor<LLFirSessionConfigurator>(
+        "org.jetbrains.kotlin.llFirSessionConfigurator",
+        LLFirSessionConfigurator::class.java
+    )
+
+    fun configure(session: LLFirSession)
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionFactory.kt

@@ -50,8 +50,7 @@ internal object LLFirSessionFactory {
         globalResolveComponents: LLFirGlobalResolveComponents,
         sessionInvalidator: LLFirSessionInvalidator,
         sessionsCache: MutableMap<KtModule, LLFirSession>,
-        librariesSessionFactory: LLFirLibrarySessionFactory,
-        configureSession: (LLFirSession.() -> Unit)? = null
+        librariesSessionFactory: LLFirLibrarySessionFactory
     ): LLFirSourcesSession {
         sessionsCache[module]?.let { return it as LLFirSourcesSession }
         checkCanceled()
@@ -110,8 +109,7 @@ internal object LLFirSessionFactory {
                     sessionsCache,
                     globalResolveComponents,
                     sessionInvalidator,
-                    librariesSessionFactory,
-                    configureSession
+                    librariesSessionFactory
                 )
 
                 add(builtinsSession)
@@ -138,7 +136,8 @@ internal object LLFirSessionFactory {
             register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
             register(LLFirFirClassByPsiClassProvider::class, LLFirFirClassByPsiClassProvider(this))
 
-            configureSession?.invoke(this)
+            LLFirSessionConfigurator.configure(this)
+
             extensionService.additionalCheckers.forEach(session.checkersComponent::register)
         }
     }
@@ -151,8 +150,7 @@ internal object LLFirSessionFactory {
         sessionInvalidator: LLFirSessionInvalidator,
         builtinSession: LLFirBuiltinsAndCloneableSession,
         sessionsCache: MutableMap<KtModule, LLFirSession>,
-        languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT,
-        configureSession: (LLFirSession.() -> Unit)? = null
+        languageVersionSettings: LanguageVersionSettings = LanguageVersionSettingsImpl.DEFAULT
     ): LLFirLibraryOrLibrarySourceResolvableModuleSession {
         LLFirLibraryOrLibrarySourceResolvableModuleSession.checkIsValidKtModule(module)
         sessionsCache[module]?.let { return it as LLFirLibraryOrLibrarySourceResolvableModuleSession }
@@ -237,7 +235,7 @@ internal object LLFirSessionFactory {
             register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
             register(LLFirFirClassByPsiClassProvider::class, LLFirFirClassByPsiClassProvider(this))
 
-            configureSession?.invoke(this)
+            LLFirSessionConfigurator.configure(this)
         }
     }
 
@@ -246,8 +244,7 @@ internal object LLFirSessionFactory {
         module: KtScriptModule,
         sessionInvalidator: LLFirSessionInvalidator,
         sessionsCache: MutableMap<KtModule, LLFirSession>,
-        librariesSessionFactory: LLFirLibrarySessionFactory,
-        configureSession: (LLFirSession.() -> Unit)? = null
+        librariesSessionFactory: LLFirLibrarySessionFactory
     ): LLFirScriptSession {
         sessionsCache[module]?.let { return it as LLFirScriptSession }
         checkCanceled()
@@ -294,8 +291,7 @@ internal object LLFirSessionFactory {
                     sessionsCache,
                     globalResolveComponents,
                     sessionInvalidator,
-                    librariesSessionFactory,
-                    configureSession
+                    librariesSessionFactory
                 )
 
                 add(builtinsSession)
@@ -321,7 +317,7 @@ internal object LLFirSessionFactory {
             register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
             register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.Empty)
 
-            configureSession?.invoke(this)
+            LLFirSessionConfigurator.configure(this)
         }
     }
 
@@ -329,8 +325,7 @@ internal object LLFirSessionFactory {
         project: Project,
         module: KtNotUnderContentRootModule,
         sessionInvalidator: LLFirSessionInvalidator,
-        sessionsCache: MutableMap<KtModule, LLFirSession>,
-        configureSession: (LLFirSession.() -> Unit)? = null
+        sessionsCache: MutableMap<KtModule, LLFirSession>
     ): LLFirNonUnderContentRootResolvableModuleSession {
         sessionsCache[module]?.let { return it as LLFirNonUnderContentRootResolvableModuleSession }
         checkCanceled()
@@ -391,7 +386,7 @@ internal object LLFirSessionFactory {
             register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
             register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotations.Empty)
 
-            configureSession?.invoke(this)
+            LLFirSessionConfigurator.configure(this)
         }
     }
 
@@ -418,8 +413,7 @@ internal object LLFirSessionFactory {
         sessionsCache: MutableMap<KtModule, LLFirSession>,
         globalResolveComponents: LLFirGlobalResolveComponents,
         sessionInvalidator: LLFirSessionInvalidator,
-        librariesSessionFactory: LLFirLibrarySessionFactory,
-        configureSession: (LLFirSession.() -> Unit)?
+        librariesSessionFactory: LLFirLibrarySessionFactory
     ) {
         val project = module.project
 
@@ -440,8 +434,7 @@ internal object LLFirSessionFactory {
                     globalResolveComponents,
                     sessionInvalidator,
                     sessionsCache,
-                    librariesSessionFactory = librariesSessionFactory,
-                    configureSession = configureSession,
+                    librariesSessionFactory = librariesSessionFactory
                 )
             }
 
@@ -466,6 +459,13 @@ internal object LLFirSessionFactory {
     }
 }
 
+internal fun LLFirSessionConfigurator.Companion.configure(session: LLFirSession) {
+    val project = session.project
+    for (extension in extensionPointName.getExtensionList(project)) {
+        extension.configure(session)
+    }
+}
+
 @Deprecated(
     "This is a dirty hack used only for one usage (building fir for psi from stubs) and it should be removed after fix of that usage",
     level = DeprecationLevel.ERROR

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionProviderStorage.kt

@@ -24,33 +24,27 @@ class LLFirSessionProviderStorage(val project: Project) {
 
     private val globalComponents = LLFirGlobalResolveComponents(project)
 
-    fun getSessionProvider(
-        useSiteKtModule: KtModule,
-        configureSession: (LLFirSession.() -> Unit)? = null
-    ): LLFirSessionProvider = when (useSiteKtModule) {
+    fun getSessionProvider(useSiteKtModule: KtModule): LLFirSessionProvider = when (useSiteKtModule) {
         is KtSourceModule -> {
-            createSessionProviderForSourceSession(useSiteKtModule, configureSession)
+            createSessionProviderForSourceSession(useSiteKtModule)
         }
 
         is KtLibraryModule, is KtLibrarySourceModule -> {
-            createSessionProviderForLibraryOrLibrarySource(useSiteKtModule, configureSession)
+            createSessionProviderForLibraryOrLibrarySource(useSiteKtModule)
         }
 
         is KtScriptModule -> {
-            createSessionProviderForScriptSession(useSiteKtModule, configureSession)
+            createSessionProviderForScriptSession(useSiteKtModule)
         }
 
         is KtNotUnderContentRootModule -> {
-            createSessionProviderForNotUnderContentRootSession(useSiteKtModule, configureSession)
+            createSessionProviderForNotUnderContentRootSession(useSiteKtModule)
         }
 
         else -> error("Unexpected ${useSiteKtModule::class.simpleName}")
     }
 
-    private fun createSessionProviderForSourceSession(
-        useSiteKtModule: KtSourceModule,
-        configureSession: (LLFirSession.() -> Unit)?
-    ): LLFirSessionProvider {
+    private fun createSessionProviderForSourceSession(useSiteKtModule: KtSourceModule): LLFirSessionProvider {
         val (sessions, session) = sourceAsUseSiteSessionCache.withMappings { mappings ->
             val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
             val session = LLFirSessionFactory.createSourcesSession(
@@ -59,8 +53,7 @@ class LLFirSessionProviderStorage(val project: Project) {
                 globalComponents,
                 sourceAsUseSiteSessionCache.sessionInvalidator,
                 sessions,
-                librariesSessionFactory,
-                configureSession = configureSession,
+                librariesSessionFactory
             )
             sessions to session
         }
@@ -68,10 +61,7 @@ class LLFirSessionProviderStorage(val project: Project) {
     }
 
 
-    private fun createSessionProviderForLibraryOrLibrarySource(
-        useSiteKtModule: KtModule,
-        configureSession: (LLFirSession.() -> Unit)?
-    ): LLFirSessionProvider {
+    private fun createSessionProviderForLibraryOrLibrarySource(useSiteKtModule: KtModule): LLFirSessionProvider {
         val (sessions, session) = libraryAsUseSiteSessionCache.withMappings { mappings ->
             val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
             val session = LLFirSessionFactory.createLibraryOrLibrarySourceResolvableSession(
@@ -80,18 +70,14 @@ class LLFirSessionProviderStorage(val project: Project) {
                 globalComponents,
                 libraryAsUseSiteSessionCache.sessionInvalidator,
                 builtInsSessionFactory.getBuiltinsSession(useSiteKtModule.platform),
-                sessions,
-                configureSession = configureSession,
+                sessions
             )
             sessions to session
         }
         return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
     }
 
-    private fun createSessionProviderForScriptSession(
-        useSiteKtModule: KtScriptModule,
-        configureSession: (LLFirSession.() -> Unit)?
-    ): LLFirSessionProvider {
+    private fun createSessionProviderForScriptSession(useSiteKtModule: KtScriptModule): LLFirSessionProvider {
         val (sessions, session) = sourceAsUseSiteSessionCache.withMappings { mappings ->
             val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
             val session = LLFirSessionFactory.createScriptSession(
@@ -99,26 +85,21 @@ class LLFirSessionProviderStorage(val project: Project) {
                 useSiteKtModule,
                 sourceAsUseSiteSessionCache.sessionInvalidator,
                 sessions,
-                librariesSessionFactory,
-                configureSession = configureSession,
+                librariesSessionFactory
             )
             sessions to session
         }
         return LLFirSessionProvider(project, session, KtModuleToSessionMappingByWeakValueMapImpl(sessions))
     }
 
-    private fun createSessionProviderForNotUnderContentRootSession(
-        useSiteKtModule: KtNotUnderContentRootModule,
-        configureSession: (LLFirSession.() -> Unit)?
-    ): LLFirSessionProvider {
+    private fun createSessionProviderForNotUnderContentRootSession(useSiteKtModule: KtNotUnderContentRootModule): LLFirSessionProvider {
         val (sessions, session) = notUnderContentRootSessionCache.withMappings { mappings ->
             val sessions = mutableMapOf<KtModule, LLFirSession>().apply { putAll(mappings) }
             val session = LLFirSessionFactory.createNotUnderContentRootResolvableSession(
                 project,
                 useSiteKtModule,
                 notUnderContentRootSessionCache.sessionInvalidator,
-                sessions,
-                configureSession = configureSession,
+                sessions
             )
             sessions to session
         }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/AbstractDiagnosticTraversalCounterTest.kt

@@ -11,11 +11,16 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDiagnosticsFor
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.BeforeElementDiagnosticCollectionHandler
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
+import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.beforeElementDiagnosticCollectionHandler
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.fir.PersistentCheckerContextFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.renderWithClassName
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolveWithClearCaches
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.useFirSessionConfigurator
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.analysis.collectors.AbstractDiagnosticCollectorVisitor
@@ -26,6 +31,7 @@ import org.jetbrains.kotlin.fir.resolve.SessionHolderImpl
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.services.TestModuleStructure
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.services.assertions
@@ -36,21 +42,21 @@ import org.jetbrains.kotlin.test.services.assertions
 abstract class AbstractDiagnosticTraversalCounterTest : AbstractLowLevelApiSingleFileTest() {
     override val configurator = AnalysisApiFirSourceTestConfigurator(analyseInDependentSession = false)
 
+    override fun configureTest(builder: TestConfigurationBuilder) {
+        super.configureTest(builder)
+        builder.apply {
+            useFirSessionConfigurator { BeforeElementLLFirSessionConfigurator() }
+        }
+    }
+
     override fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices) {
-        val handler = BeforeElementTestDiagnosticCollectionHandler()
-        resolveWithClearCaches(
-            ktFile,
-            configureSession = {
-                @OptIn(SessionConfiguration::class)
-                register(BeforeElementDiagnosticCollectionHandler::class, handler)
-            }
-        ) { firResolveSession ->
+        resolveWithClearCaches(ktFile) { firResolveSession ->
             // we should get diagnostics before we resolve the whole file by  ktFile.getOrBuildFir
             ktFile.collectDiagnosticsForFile(firResolveSession, DiagnosticCheckerFilter.ONLY_COMMON_CHECKERS)
 
             val firFile = ktFile.getOrBuildFirOfType<FirFile>(firResolveSession)
 
-            val errorElements = collectErrorElements(firFile, handler)
+            val errorElements = collectErrorElements(firFile)
 
             if (errorElements.isNotEmpty()) {
                 val zeroElements = errorElements.filter { it.second == 0 }
@@ -77,10 +83,8 @@ abstract class AbstractDiagnosticTraversalCounterTest : AbstractLowLevelApiSingl
         }
     }
 
-    private fun collectErrorElements(
-        firFile: FirElement,
-        handler: BeforeElementTestDiagnosticCollectionHandler
-    ): List<Pair<FirElement, Int>> {
+    private fun collectErrorElements(firFile: FirFile): List<Pair<FirElement, Int>> {
+        val handler = firFile.llFirSession.beforeElementDiagnosticCollectionHandler as BeforeElementTestDiagnosticCollectionHandler
         val errorElements = mutableListOf<Pair<FirElement, Int>>()
         val nonDuplicatingElements = findNonDuplicatingFirElements(firFile).filter { element ->
             when {
@@ -128,6 +132,13 @@ abstract class AbstractDiagnosticTraversalCounterTest : AbstractLowLevelApiSingl
         return elementUsageCount.filterValues { it == 1 }.keys
     }
 
+    private class BeforeElementLLFirSessionConfigurator : LLFirSessionConfigurator {
+        @OptIn(SessionConfiguration::class)
+        override fun configure(session: LLFirSession) {
+            val handler = BeforeElementTestDiagnosticCollectionHandler()
+            session.register(BeforeElementDiagnosticCollectionHandler::class, handler)
+        }
+    }
 
     class BeforeElementTestDiagnosticCollectionHandler : BeforeElementDiagnosticCollectionHandler() {
         val visitedTimes = mutableMapOf<FirElement, Int>()

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/AbstractFirContextCollectionTest.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilt
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getDiagnostics
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.BeforeElementDiagnosticCollectionHandler
+import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.beforeElementDiagnosticCollectionHandler
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.fir.PersistenceContextCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.DanglingTopLevelModifierListStructureElement
@@ -26,29 +27,32 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.isSourceSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.name
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolveWithClearCaches
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirOutOfContentRootTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
+import org.jetbrains.kotlin.analysis.low.level.api.fir.useFirSessionConfigurator
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.test.services.AssertionsService
-import org.jetbrains.kotlin.test.services.TestModuleStructure
-import org.jetbrains.kotlin.test.services.TestServices
-import org.jetbrains.kotlin.test.services.assertions
+import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
+import org.jetbrains.kotlin.test.services.*
 
 abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileTest() {
+    override fun configureTest(builder: TestConfigurationBuilder) {
+        super.configureTest(builder)
+        builder.apply {
+            useFirSessionConfigurator(::BeforeElementLLFirSessionConfigurator)
+        }
+    }
+
     override fun doTestByFileStructure(ktFile: KtFile, moduleStructure: TestModuleStructure, testServices: TestServices) {
-        val handler = BeforeElementTestDiagnosticCollectionHandler(testServices.assertions)
-        resolveWithClearCaches(
-            ktFile,
-            configureSession = {
-                @OptIn(SessionConfiguration::class)
-                register(BeforeElementDiagnosticCollectionHandler::class, handler)
-            }
-        ) { firResolveSession ->
+        resolveWithClearCaches(ktFile) { firResolveSession ->
             check(firResolveSession.isSourceSession)
 
             val session = firResolveSession.getSessionFor(ktFile.getKtModule()) as LLFirResolvableModuleSession
+            val handler = session.beforeElementDiagnosticCollectionHandler as BeforeElementTestDiagnosticCollectionHandler
+
             val fileStructureCache = session.moduleComponents.fileStructureCache
 
             val fileStructure = fileStructureCache.getFileStructure(ktFile)
@@ -68,11 +72,17 @@ abstract class AbstractFirContextCollectionTest : AbstractLowLevelApiSingleFileT
         is DanglingTopLevelModifierListStructureElement -> fir
     }
 
+    private class BeforeElementLLFirSessionConfigurator(private val testServices: TestServices) : LLFirSessionConfigurator {
+        @OptIn(SessionConfiguration::class)
+        override fun configure(session: LLFirSession) {
+            val handler = BeforeElementTestDiagnosticCollectionHandler(testServices.assertions)
+            session.register(BeforeElementDiagnosticCollectionHandler::class, handler)
+        }
+    }
 
     private class BeforeElementTestDiagnosticCollectionHandler(
         private val assertions: AssertionsService
-    ) :
-        BeforeElementDiagnosticCollectionHandler() {
+    ) : BeforeElementDiagnosticCollectionHandler() {
         lateinit var elementsToCheckContext: List<FirDeclaration>
         lateinit var firFile: FirFile
 

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/firTestUtils.kt

@@ -6,20 +6,24 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirNotUnderContentRootResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirSourceResolveSession
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
+import org.jetbrains.kotlin.analysis.test.framework.services.environmentManager
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.render
-import org.jetbrains.kotlin.fir.renderer.FirRenderer
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousInitializerSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFileSymbol
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
+import org.jetbrains.kotlin.test.services.PreAnalysisHandler
+import org.jetbrains.kotlin.test.services.TestModuleStructure
+import org.jetbrains.kotlin.test.services.TestServices
 
 
 internal fun FirElement.renderWithClassName(): String =
@@ -36,13 +40,9 @@ internal fun FirBasedSymbol<*>.name(): String = when (this) {
 
 internal fun FirDeclaration.name(): String = symbol.name()
 
-internal inline fun <R> resolveWithClearCaches(
-    context: KtElement,
-    noinline configureSession: LLFirSession.() -> Unit = {},
-    action: (LLFirResolveSession) -> R,
-): R {
+internal inline fun <R> resolveWithClearCaches(context: KtElement, action: (LLFirResolveSession) -> R): R {
     val project = context.project
-    val firResolveSession = createFirResolveSessionForNoCaching(context.getKtModule(project), project, configureSession)
+    val firResolveSession = createFirResolveSessionForNoCaching(context.getKtModule(project), project)
     return action(firResolveSession)
 }
 
@@ -52,4 +52,15 @@ internal val LLFirResolveSession.isSourceSession: Boolean
             is LLFirSourceResolveSession, is LLFirNotUnderContentRootResolveSession -> true
             else -> false
         }
-    }
+    }
+
+internal fun TestConfigurationBuilder.useFirSessionConfigurator(configurator: (TestServices) -> LLFirSessionConfigurator) {
+    class ConfiguratorPreAnalysisHandler(testServices: TestServices) : PreAnalysisHandler(testServices) {
+        override fun preprocessModuleStructure(moduleStructure: TestModuleStructure) {
+            val project = testServices.environmentManager.getProject()
+            LLFirSessionConfigurator.registerExtension(project, configurator(testServices))
+        }
+    }
+
+    usePreAnalysisHandlers(::ConfiguratorPreAnalysisHandler)
+}

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/AnalysisApiFirTestServiceRegistrar.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirLi
 import org.jetbrains.kotlin.analysis.low.level.api.fir.services.KtCompilerPluginsProviderForTests
 import org.jetbrains.kotlin.analysis.low.level.api.fir.services.LLFirSealedClassInheritorsProcessorFactoryForTests
 import org.jetbrains.kotlin.analysis.low.level.api.fir.services.PackagePartProviderTestImpl
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
 import org.jetbrains.kotlin.analysis.project.structure.KtCompilerPluginsProvider
 import org.jetbrains.kotlin.analysis.providers.PackagePartProviderFactory
 import org.jetbrains.kotlin.analysis.test.framework.test.configurators.AnalysisApiTestServiceRegistrar
@@ -39,6 +40,7 @@ object AnalysisApiFirTestServiceRegistrar : AnalysisApiTestServiceRegistrar() {
     override fun registerProjectExtensionPoints(project: MockProject, testServices: TestServices) {
         IrGenerationExtension.registerExtensionPoint(project)
         FirExtensionRegistrarAdapter.registerExtensionPoint(project)
+        LLFirSessionConfigurator.registerExtensionPoint(project)
     }
 
     @OptIn(KtAnalysisApiInternals::class)