FIR IDE: add ability to run multimodule tests

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/TestConfiguration.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.test.directives.model.RegisteredDirectives
 import org.jetbrains.kotlin.test.model.*
 import org.jetbrains.kotlin.test.services.MetaTestConfigurator
 import org.jetbrains.kotlin.test.services.ModuleStructureExtractor
+import org.jetbrains.kotlin.test.services.PreAnalysisHandler
 import org.jetbrains.kotlin.test.services.TestServices
 
 typealias Constructor<T> = (TestServices) -> T
@@ -26,6 +27,8 @@ abstract class TestConfiguration {
 
     abstract val moduleStructureExtractor: ModuleStructureExtractor
 
+    abstract val preAnalysisHandlers: List<PreAnalysisHandler>
+
     abstract val metaTestConfigurators: List<MetaTestConfigurator>
 
     abstract val afterAnalysisCheckers: List<AfterAnalysisChecker>

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/TestRunner.kt

@@ -52,6 +52,7 @@ class TestRunner(private val testConfiguration: TestConfiguration) {
             throw exception
         }
 
+
         testConfiguration.metaTestConfigurators.forEach {
             if (it.shouldSkipTest()) return
         }
@@ -62,6 +63,11 @@ class TestRunner(private val testConfiguration: TestConfiguration) {
         val modules = moduleStructure.modules
         val dependencyProvider = DependencyProviderImpl(services, modules)
         services.registerDependencyProvider(dependencyProvider)
+
+        testConfiguration.preAnalysisHandlers.forEach { preprocessor ->
+            preprocessor.preprocessModuleStructure(moduleStructure)
+        }
+
         var failedException: WrappedException? = null
         try {
             for (module in modules) {

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/services/PreAnalysisHandler.kt

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2010-2021 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.test.services
+
+abstract class PreAnalysisHandler(protected val testServices: TestServices) {
+    abstract fun preprocessModuleStructure(moduleStructure: TestModuleStructure)
+}

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/builders/TestConfigurationBuilder.kt

@@ -28,6 +28,7 @@ class TestConfigurationBuilder {
     private val sourcePreprocessors: MutableList<Constructor<SourceFilePreprocessor>> = mutableListOf()
     private val additionalMetaInfoProcessors: MutableList<Constructor<AdditionalMetaInfoProcessor>> = mutableListOf()
     private val environmentConfigurators: MutableList<Constructor<EnvironmentConfigurator>> = mutableListOf()
+    private val preAnalysisHandlers: MutableList<Constructor<PreAnalysisHandler>> = mutableListOf()
 
     private val additionalSourceProviders: MutableList<Constructor<AdditionalSourceProvider>> = mutableListOf()
     private val moduleStructureTransformers: MutableList<ModuleStructureTransformer> = mutableListOf()
@@ -131,6 +132,10 @@ class TestConfigurationBuilder {
         this.environmentConfigurators += environmentConfigurators
     }
 
+    fun usePreAnalysisHandlers(vararg handlers: Constructor<PreAnalysisHandler>) {
+        this.preAnalysisHandlers += handlers
+    }
+
     fun useMetaInfoProcessors(vararg updaters: Constructor<AdditionalMetaInfoProcessor>) {
         additionalMetaInfoProcessors += updaters
     }
@@ -189,6 +194,7 @@ class TestConfigurationBuilder {
             additionalMetaInfoProcessors,
             environmentConfigurators,
             additionalSourceProviders,
+            preAnalysisHandlers,
             moduleStructureTransformers,
             metaTestConfigurators,
             afterAnalysisCheckers,

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/impl/TestConfigurationImpl.kt

@@ -33,6 +33,7 @@ class TestConfigurationImpl(
     environmentConfigurators: List<Constructor<EnvironmentConfigurator>>,
 
     additionalSourceProviders: List<Constructor<AdditionalSourceProvider>>,
+    preAnalysisHandlers: List<Constructor<PreAnalysisHandler>>,
     moduleStructureTransformers: List<ModuleStructureTransformer>,
     metaTestConfigurators: List<Constructor<MetaTestConfigurator>>,
     afterAnalysisCheckers: List<Constructor<AfterAnalysisChecker>>,
@@ -70,6 +71,9 @@ class TestConfigurationImpl(
             }
         }
 
+    override val preAnalysisHandlers: List<PreAnalysisHandler> =
+        preAnalysisHandlers.map { it.invoke(testServices) }
+
     override val moduleStructureExtractor: ModuleStructureExtractor = ModuleStructureExtractorImpl(
         testServices,
         additionalSourceProviders.map { it.invoke(testServices) }.also {

idea/idea-frontend-api/build.gradle.kts

@@ -3,6 +3,10 @@ plugins {
     id("jps-compatible")
 }
 
+kotlin {
+    explicitApiWarning()
+}
+
 dependencies {
     compileOnly(project(":kotlin-reflect-api"))
 

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

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.analyzer.ModuleSourceInfoBase
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName

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

@@ -8,6 +8,8 @@ package org.jetbrains.kotlin.idea.fir.low.level.api
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.analyzer.ModuleSourceInfoBase
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform

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

@@ -23,7 +23,7 @@ import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.psi.KtElement
 
 abstract class FirModuleResolveStateConfigurator {
-    abstract fun createPackagePartsProvider(scope: GlobalSearchScope): PackagePartProvider
+    abstract fun createPackagePartsProvider(moduleInfo: ModuleSourceInfoBase, scope: GlobalSearchScope): PackagePartProvider
 
     abstract fun createModuleDataProvider(moduleInfo: ModuleSourceInfoBase): ModuleDataProvider
 

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

@@ -176,7 +176,7 @@ internal object FirIdeSessionFactory {
             setProjectInstance(project)
             setScope(searchScope)
         }
-        val packagePartProvider = project.stateConfigurator.createPackagePartsProvider(searchScope)
+        val packagePartProvider = project.stateConfigurator.createPackagePartsProvider(mainModuleInfo, searchScope)
 
         val kotlinClassFinder = VirtualFileFinderFactory.getInstance(project).create(searchScope)
         FirIdeLibrariesSession(project, searchScope, builtinTypes).apply session@{

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/compiler/based/AbstractCompilerBasedTest.kt

@@ -25,20 +25,19 @@ import org.jetbrains.kotlin.fir.DependencyListForCliModule
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.SealedClassInheritorsProvider
 import org.jetbrains.kotlin.fir.deserialization.ModuleDataProvider
+import org.jetbrains.kotlin.fir.java.FirJavaElementFinder
 import org.jetbrains.kotlin.fir.session.FirModuleInfoBasedModuleData
 import org.jetbrains.kotlin.idea.asJava.IDEKotlinAsJavaFirSupport
 import org.jetbrains.kotlin.idea.fir.low.level.api.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.DiagnosticCheckerFilter
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveStateConfigurator
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.KotlinOutOfBlockModificationTrackerFactory
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.*
-import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.AbstractLowLevelApiTest.Companion.reRegisterJavaElementFinder
 import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.DeclarationProviderTestImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.KotlinOutOfBlockModificationTrackerFactoryTestImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.KtPackageProviderTestImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.test.base.SealedClassInheritorsProviderTestImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.transformers.FirLazyTransformerForIDE
+import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProvider
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSessionProvider
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.TargetPlatform
@@ -47,14 +46,12 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
 import org.jetbrains.kotlin.test.TestConfiguration
+import org.jetbrains.kotlin.test.bind
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.builders.testConfiguration
 import org.jetbrains.kotlin.test.directives.FirDiagnosticsDirectives
 import org.jetbrains.kotlin.test.directives.model.DirectivesContainer
-import org.jetbrains.kotlin.test.frontend.fir.FirModuleInfoProvider
+import org.jetbrains.kotlin.test.frontend.fir.*
-import org.jetbrains.kotlin.test.frontend.fir.FirOutputArtifact
-import org.jetbrains.kotlin.test.frontend.fir.firModuleInfoProvider
-import org.jetbrains.kotlin.test.frontend.fir.getAnalyzerServices
 import org.jetbrains.kotlin.test.model.*
 import org.jetbrains.kotlin.test.runners.AbstractKotlinCompilerTest
 import org.jetbrains.kotlin.test.services.*
@@ -87,6 +84,8 @@ abstract class AbstractCompilerBasedTest : AbstractKotlinCompilerTest() {
         configureTest()
         defaultConfiguration(this)
         unregisterAllFacades()
+        useAdditionalService(::TestModuleInfoProvider)
+        usePreAnalysisHandlers(::ModuleRegistrarPreAnalysisHandler.bind(disposable))
         useFrontendFacades(::LowLevelFirFrontendFacade)
     }
 
@@ -96,30 +95,18 @@ abstract class AbstractCompilerBasedTest : AbstractKotlinCompilerTest() {
     inner class LowLevelFirFrontendFacade(
         testServices: TestServices
     ) : FrontendFacade<FirOutputArtifact>(testServices, FrontendKinds.FIR) {
-        override val additionalServices: List<ServiceRegistrationData>
-            get() = listOf(service(::FirModuleInfoProvider))
 
         override val additionalDirectives: List<DirectivesContainer>
             get() = listOf(FirDiagnosticsDirectives)
 
         override fun analyze(module: TestModule): FirOutputArtifact {
+            val moduleInfoProvider = testServices.moduleInfoProvider
+            val moduleInfo = moduleInfoProvider.getModuleInfo(module.name)
+
             val project = testServices.compilerConfigurationProvider.getProject(module)
-            val ktFiles = testServices.sourceFileProvider.getKtFilesForSourceFiles(module.files, project)
-
-            reRegisterJavaElementFinder(project)
-
-            val moduleInfo = TestModuleInfo(module)
-            testServices.firModuleInfoProvider.registerModuleData(module, FirModuleInfoBasedModuleData(moduleInfo))
-            val configurator =
-                FirModuleResolveStateConfiguratorForSingleModuleTestImpl(testServices, module, ktFiles, moduleInfo, disposable)
-
-            with(project as MockProject) {
-                registerTestServices(configurator, ktFiles)
-            }
-
             val resolveState = createResolveStateForNoCaching(moduleInfo, project)
 
-            val allFirFiles = ktFiles.map { (testFile, psiFile) ->
+            val allFirFiles = moduleInfo.testFilesToKtFiles.map { (testFile, psiFile) ->
                 testFile to psiFile.getOrBuildFirFile(resolveState)
             }.toMap()
 
@@ -149,12 +136,7 @@ abstract class AbstractCompilerBasedTest : AbstractKotlinCompilerTest() {
     private fun ignoreTest(filePath: String, configuration: TestConfiguration): Boolean {
         val modules = configuration.moduleStructureExtractor.splitTestDataByModules(filePath, configuration.directives)
 
-        if (modules.modules.size > 1) {
+        if (modules.modules.none { it.files.any { it.isKtFile } }) {
-            return true // multimodule tests are not supported
-        }
-
-        val singleModule = modules.modules.single()
-        if (singleModule.files.none { it.isKtFile }) {
             return true // nothing to highlight
         }
 
@@ -162,15 +144,80 @@ abstract class AbstractCompilerBasedTest : AbstractKotlinCompilerTest() {
     }
 }
 
-class TestModuleInfo(
-    val testModule: TestModule,
-) : ModuleInfo, ModuleSourceInfoBase {
-    override val name: Name get() = Name.identifier(testModule.name)
 
-    override fun dependencies(): List<ModuleInfo> {
+class TestModuleInfoProvider(private val testServices: TestServices) : TestService {
-        return emptyList()
+    private val cache = mutableMapOf<String, TestModuleSourceInfo>()
+
+    fun registerModuleInfo(testModule: TestModule, ktFiles: Map<TestFile, KtFile>) {
+        cache[testModule.name] = TestModuleSourceInfo(testModule, ktFiles, testServices)
     }
 
+    fun getModuleInfoByKtFile(ktFile: KtFile): TestModuleSourceInfo =
+        cache.values.first { moduleSourceInfo ->
+            ktFile in moduleSourceInfo.ktFiles
+        }
+
+    fun getModuleInfo(moduleName: String): TestModuleSourceInfo =
+        cache.getValue(moduleName)
+}
+
+val TestServices.moduleInfoProvider: TestModuleInfoProvider by TestServices.testServiceAccessor()
+
+class ModuleRegistrarPreAnalysisHandler(
+    testServices: TestServices,
+    private val parentDisposable: Disposable
+) : PreAnalysisHandler(testServices) {
+
+    private val moduleInfoProvider = testServices.moduleInfoProvider
+
+    override fun preprocessModuleStructure(moduleStructure: TestModuleStructure) {
+        // todo rework after all modules will have the same Project instance
+        val ktFilesByModule = moduleStructure.modules.associateWith { testModule ->
+            val project = testServices.compilerConfigurationProvider.getProject(testModule)
+            testServices.sourceFileProvider.getKtFilesForSourceFiles(testModule.files, project)
+        }
+
+        val allKtFiles = ktFilesByModule.values.flatMap { it.values.toList() }
+
+        ktFilesByModule.forEach { (testModule, ktFiles) ->
+            val project = testServices.compilerConfigurationProvider.getProject(testModule)
+            moduleInfoProvider.registerModuleInfo(testModule, ktFiles)
+
+            val configurator = FirModuleResolveStateConfiguratorForSingleModuleTestImpl(
+                project,
+                testServices,
+                parentDisposable
+            )
+            (project as MockProject).registerTestServices(configurator, allKtFiles)
+        }
+    }
+}
+
+class TestModuleSourceInfo(
+    val testModule: TestModule,
+    val testFilesToKtFiles: Map<TestFile, KtFile>,
+    testServices: TestServices,
+) : ModuleInfo, ModuleSourceInfoBase {
+    private val moduleInfoProvider = testServices.moduleInfoProvider
+
+    val ktFiles = testFilesToKtFiles.values.toSet()
+    val moduleData = FirModuleInfoBasedModuleData(this)
+
+    override val name: Name
+        get() = Name.identifier(testModule.name)
+
+    override fun dependencies(): List<ModuleInfo> =
+        testModule.regularDependencies
+            .map { moduleInfoProvider.getModuleInfo(it.moduleName) }
+
+    override val expectedBy: List<ModuleInfo>
+        get() = testModule.dependsOnDependencies
+            .map { moduleInfoProvider.getModuleInfo(it.moduleName) }
+
+    override fun modulesWhoseInternalsAreVisible(): Collection<ModuleInfo> =
+        testModule.friendDependencies
+            .map { moduleInfoProvider.getModuleInfo(it.moduleName) }
+
     override val platform: TargetPlatform
         get() = testModule.targetPlatform
 
@@ -180,32 +227,25 @@ class TestModuleInfo(
 
 
 class FirModuleResolveStateConfiguratorForSingleModuleTestImpl(
-    private val testServices: TestServices,
+    private val project: Project,
-    private val testModule: TestModule,
+    testServices: TestServices,
-    private val ktFiles: Map<TestFile, KtFile>,
-    private val moduleInfo: ModuleInfo,
     private val parentDisposable: Disposable,
 ) : FirModuleResolveStateConfigurator() {
-    val moduleInfoProvider = testServices.firModuleInfoProvider
+    private val moduleInfoProvider = testServices.moduleInfoProvider
-    val compilerConfigurationProvider = testServices.compilerConfigurationProvider
+    private val compilerConfigurationProvider = testServices.compilerConfigurationProvider
-
+    private val librariesScope = ProjectScope.getLibrariesScope(project)//todo incorrect?
-    val project = compilerConfigurationProvider.getProject(testModule)
-
-    val languageVersionSettings = testModule.languageVersionSettings
-    val packagePartProviderFactory = compilerConfigurationProvider.getPackagePartProviderFactory(testModule)
-
-    val configuration = compilerConfigurationProvider.getCompilerConfiguration(testModule)
-
-    val librariesScope = ProjectScope.getLibrariesScope(project)
-    val sourcesScope = TopDownAnalyzerFacadeForJVM.newModuleSearchScope(project, ktFiles.values)
-
     private val sealedClassInheritorsProvider = SealedClassInheritorsProviderTestImpl()
 
-    override fun createPackagePartsProvider(scope: GlobalSearchScope): PackagePartProvider {
+    override fun createPackagePartsProvider(moduleInfo: ModuleSourceInfoBase, scope: GlobalSearchScope): PackagePartProvider {
-        return packagePartProviderFactory.invoke(scope)
+        require(moduleInfo is TestModuleSourceInfo)
+        val factory = compilerConfigurationProvider.getPackagePartProviderFactory(moduleInfo.testModule)
+        return factory(scope)
     }
 
+
     override fun createModuleDataProvider(moduleInfo: ModuleSourceInfoBase): ModuleDataProvider {
+        require(moduleInfo is TestModuleSourceInfo)
+        val configuration = compilerConfigurationProvider.getCompilerConfiguration(moduleInfo.testModule)
         return DependencyListForCliModule.build(
             moduleInfo.name,
             moduleInfo.platform,
@@ -216,18 +256,20 @@ class FirModuleResolveStateConfiguratorForSingleModuleTestImpl(
 
             friendDependencies(configuration[JVMConfigurationKeys.FRIEND_PATHS] ?: emptyList())
 
-            sourceDependencies(moduleInfoProvider.getRegularDependentSourceModules(testModule))
+            sourceDependencies(moduleInfo.moduleData.dependencies)
-            sourceFriendsDependencies(moduleInfoProvider.getDependentFriendSourceModules(testModule))
+            sourceFriendsDependencies(moduleInfo.moduleData.friendDependencies)
-            sourceDependsOnDependencies(moduleInfoProvider.getDependentDependsOnSourceModules(testModule))
+            sourceDependsOnDependencies(moduleInfo.moduleData.dependsOnDependencies)
         }.moduleDataProvider
     }
 
     override fun getLanguageVersionSettings(moduleInfo: ModuleSourceInfoBase): LanguageVersionSettings {
-        return languageVersionSettings
+        require(moduleInfo is TestModuleSourceInfo)
+        return moduleInfo.testModule.languageVersionSettings
     }
 
     override fun getModuleSourceScope(moduleInfo: ModuleSourceInfoBase): GlobalSearchScope {
-        return sourcesScope
+        require(moduleInfo is TestModuleSourceInfo)
+        return TopDownAnalyzerFacadeForJVM.newModuleSearchScope(project, moduleInfo.testFilesToKtFiles.values)
     }
 
     override fun createScopeForModuleLibraries(moduleInfo: ModuleSourceInfoBase): GlobalSearchScope {
@@ -239,7 +281,8 @@ class FirModuleResolveStateConfiguratorForSingleModuleTestImpl(
     }
 
     override fun getModuleInfoFor(element: KtElement): ModuleInfo {
-        return moduleInfo
+        val containingFile = element.containingKtFile
+        return moduleInfoProvider.getModuleInfoByKtFile(containingFile)
     }
 
     override fun configureSourceSession(session: FirSession) {
@@ -248,7 +291,7 @@ class FirModuleResolveStateConfiguratorForSingleModuleTestImpl(
 
 fun MockProject.registerTestServices(
     configurator: FirModuleResolveStateConfiguratorForSingleModuleTestImpl,
-    ktFiles: Map<TestFile, KtFile>
+    allKtFiles: List<KtFile>,
 ) {
     registerService(
         FirModuleResolveStateConfigurator::class.java,
@@ -261,12 +304,30 @@ fun MockProject.registerTestServices(
     )
     registerService(KtDeclarationProviderFactory::class.java, object : KtDeclarationProviderFactory() {
         override fun createDeclarationProvider(searchScope: GlobalSearchScope): DeclarationProvider {
-            return DeclarationProviderTestImpl(searchScope, ktFiles.values)
+            return DeclarationProviderTestImpl(searchScope, allKtFiles.filter { searchScope.contains(it.virtualFile) })
         }
     })
     registerService(KtPackageProviderFactory::class.java, object : KtPackageProviderFactory() {
         override fun createPackageProvider(searchScope: GlobalSearchScope): KtPackageProvider {
-            return KtPackageProviderTestImpl(searchScope, ktFiles.values)
+            return KtPackageProviderTestImpl(searchScope, allKtFiles.filter { searchScope.contains(it.virtualFile) })
         }
     })
+    reRegisterJavaElementFinder(this)
+}
+
+@OptIn(InvalidWayOfUsingAnalysisSession::class)
+private fun reRegisterJavaElementFinder(project: Project) {
+    PsiElementFinder.EP.getPoint(project).unregisterExtension(JavaElementFinder::class.java)
+    with(project as MockProject) {
+        picoContainer.registerComponentInstance(
+            KtAnalysisSessionProvider::class.qualifiedName,
+            KtFirAnalysisSessionProvider(this)
+        )
+        picoContainer.unregisterComponent(KotlinAsJavaSupport::class.qualifiedName)
+        picoContainer.registerComponentInstance(
+            KotlinAsJavaSupport::class.qualifiedName,
+            IDEKotlinAsJavaFirSupport(project)
+        )
+    }
+    PsiElementFinder.EP.getPoint(project).registerExtension(JavaElementFinder(project))
 }

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/test/base/AbstractLowLevelApiTest.kt

@@ -11,19 +11,14 @@ import com.intellij.psi.PsiElementFinder
 import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.fir.session.FirModuleInfoBasedModuleData
-import org.jetbrains.kotlin.idea.asJava.IDEKotlinAsJavaFirSupport
 import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.FirModuleResolveStateConfiguratorForSingleModuleTestImpl
-import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.TestModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.registerTestServices
-import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
+import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.*
-import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSessionProvider
-import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSessionProvider
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.test.bind
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.builders.testConfiguration
-import org.jetbrains.kotlin.test.frontend.fir.FirModuleInfoProvider
-import org.jetbrains.kotlin.test.frontend.fir.firModuleInfoProvider
 import org.jetbrains.kotlin.test.model.DependencyKind
 import org.jetbrains.kotlin.test.model.FrontendKinds
 import org.jetbrains.kotlin.test.services.*
@@ -54,8 +49,8 @@ abstract class AbstractLowLevelApiTest : TestWithDisposable() {
 
         useSourcePreprocessor(::ExpressionMarkersSourceFilePreprocessor)
         useAdditionalService(::ExpressionMarkerProvider)
-        useAdditionalService(::FirModuleInfoProvider)
+        useAdditionalService(::TestModuleInfoProvider)
-
+        usePreAnalysisHandlers(::ModuleRegistrarPreAnalysisHandler.bind(disposable))
         configureTest(this)
 
         this.testInfo = this@AbstractLowLevelApiTest.testInfo
@@ -70,25 +65,6 @@ abstract class AbstractLowLevelApiTest : TestWithDisposable() {
 
     open fun configureTest(builder: TestConfigurationBuilder) {}
 
-    companion object {
-        @OptIn(InvalidWayOfUsingAnalysisSession::class)
-        fun reRegisterJavaElementFinder(project: Project) {
-            PsiElementFinder.EP.getPoint(project).unregisterExtension(JavaElementFinder::class.java)
-            with(project as MockProject) {
-                picoContainer.registerComponentInstance(
-                    KtAnalysisSessionProvider::class.qualifiedName,
-                    KtFirAnalysisSessionProvider(this)
-                )
-                picoContainer.unregisterComponent(KotlinAsJavaSupport::class.qualifiedName)
-                picoContainer.registerComponentInstance(
-                    KotlinAsJavaSupport::class.qualifiedName,
-                    IDEKotlinAsJavaFirSupport(project)
-                )
-            }
-            PsiElementFinder.EP.getPoint(project).registerExtension(JavaElementFinder(project))
-        }
-    }
-
     protected fun runTest(path: String) {
         testDataPath = Paths.get(path)
         val testConfiguration = testConfiguration(path, configure)
@@ -96,32 +72,25 @@ abstract class AbstractLowLevelApiTest : TestWithDisposable() {
         val moduleStructure = testConfiguration.moduleStructureExtractor.splitTestDataByModules(
             path,
             testConfiguration.directives,
-        ).also {
+        ).also { testModuleStructure ->
-            testConfiguration.testServices.register(TestModuleStructure::class, it)
+            testConfiguration.testServices.register(TestModuleStructure::class, testModuleStructure)
+            testConfiguration.preAnalysisHandlers.forEach { preprocessor ->
+                preprocessor.preprocessModuleStructure(testModuleStructure)
+            }
         }
 
+
         val singleModule = moduleStructure.modules.single()
         val project = testServices.compilerConfigurationProvider.getProject(singleModule)
-        val ktFiles = testServices.sourceFileProvider.getKtFilesForSourceFiles(singleModule.files, project)
+        val moduleInfoProvider = testServices.moduleInfoProvider
 
-        reRegisterJavaElementFinder(project)
+        val moduleInfo = moduleInfoProvider.getModuleInfo(singleModule.name)
-
-        val moduleInfo = TestModuleInfo(singleModule)
-        testServices.firModuleInfoProvider.registerModuleData(singleModule, FirModuleInfoBasedModuleData(moduleInfo))
-        val configurator = FirModuleResolveStateConfiguratorForSingleModuleTestImpl(
-            testServices,
-            singleModule,
-            ktFiles,
-            moduleInfo,
-            disposable
-        )
 
         with(project as MockProject) {
-            registerTestServices(configurator, ktFiles)
             registerServicesForProject(this)
         }
 
-        doTestByFileStructure(ktFiles.values.toList(), moduleStructure, testServices)
+        doTestByFileStructure(moduleInfo.ktFiles.toList(), moduleStructure, testServices)
     }
 
     protected open fun registerServicesForProject(project: MockProject) {}