FIR IDE: add ability to run multimodule tests

2021-06-22 17:12:39 +02:00
parent 1cb34541bd
commit a9bb577154
12 changed files with 179 additions and 112 deletions
@@ -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>
@@ -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) {
@@ -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)
+}
@@ -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,
@@ -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 {
@@ -3,6 +3,10 @@ plugins {
    id("jps-compatible")
 }

+kotlin {
+    explicitApiWarning()
+}
+
 dependencies {
    compileOnly(project(":kotlin-reflect-api"))

@@ -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
@@ -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
@@ -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

@@ -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@{
@@ -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.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.api.*
 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.FirOutputArtifact
-import org.jetbrains.kotlin.test.frontend.fir.firModuleInfoProvider
-import org.jetbrains.kotlin.test.frontend.fir.getAnalyzerServices
+import org.jetbrains.kotlin.test.frontend.fir.*
 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) {
-            return true // multimodule tests are not supported
-        }
-
-        val singleModule = modules.modules.single()
-        if (singleModule.files.none { it.isKtFile }) {
+        if (modules.modules.none { it.files.any { 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> {
-        return emptyList()
+class TestModuleInfoProvider(private val testServices: TestServices) : TestService {
+    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 testModule: TestModule,
-    private val ktFiles: Map<TestFile, KtFile>,
-    private val moduleInfo: ModuleInfo,
+    private val project: Project,
+    testServices: TestServices,
    private val parentDisposable: Disposable,
 ) : FirModuleResolveStateConfigurator() {
-    val moduleInfoProvider = testServices.firModuleInfoProvider
-    val compilerConfigurationProvider = testServices.compilerConfigurationProvider
-
-    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 moduleInfoProvider = testServices.moduleInfoProvider
+    private val compilerConfigurationProvider = testServices.compilerConfigurationProvider
+    private val librariesScope = ProjectScope.getLibrariesScope(project)//todo incorrect?
    private val sealedClassInheritorsProvider = SealedClassInheritorsProviderTestImpl()

-    override fun createPackagePartsProvider(scope: GlobalSearchScope): PackagePartProvider {
-        return packagePartProviderFactory.invoke(scope)
+    override fun createPackagePartsProvider(moduleInfo: ModuleSourceInfoBase, scope: GlobalSearchScope): PackagePartProvider {
+        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))
-            sourceFriendsDependencies(moduleInfoProvider.getDependentFriendSourceModules(testModule))
-            sourceDependsOnDependencies(moduleInfoProvider.getDependentDependsOnSourceModules(testModule))
+            sourceDependencies(moduleInfo.moduleData.dependencies)
+            sourceFriendsDependencies(moduleInfo.moduleData.friendDependencies)
+            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))
+}
@@ -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.frontend.api.KtAnalysisSessionProvider
-import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSessionProvider
+import org.jetbrains.kotlin.idea.fir.low.level.api.compiler.based.*
 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 {
-            testConfiguration.testServices.register(TestModuleStructure::class, it)
+        ).also { testModuleStructure ->
+            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 = TestModuleInfo(singleModule)
-        testServices.firModuleInfoProvider.registerModuleData(singleModule, FirModuleInfoBasedModuleData(moduleInfo))
-        val configurator = FirModuleResolveStateConfiguratorForSingleModuleTestImpl(
-            testServices,
-            singleModule,
-            ktFiles,
-            moduleInfo,
-            disposable
-        )
+        val moduleInfo = moduleInfoProvider.getModuleInfo(singleModule.name)

        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) {}