[CLI] Introduce utilities for creating FirSession hierarchy in CLI for all platforms

Also support session creation and compilation for HMPP projects ^KT-56209 Fixed
2023-02-09 22:24:32 +02:00
parent 30ea4b6b53
commit 79e4df72bf
22 changed files with 783 additions and 633 deletions
@@ -0,0 +1,411 @@
+/*
+ * 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.cli.common
+
+import org.jetbrains.kotlin.KtSourceFile
+import org.jetbrains.kotlin.analyzer.common.CommonPlatformAnalyzerServices
+import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.GroupedKtSources
+import org.jetbrains.kotlin.config.*
+import org.jetbrains.kotlin.fir.DependencyListForCliModule
+import org.jetbrains.kotlin.fir.FirModuleData
+import org.jetbrains.kotlin.fir.FirModuleDataImpl
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
+import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
+import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.fir.session.*
+import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
+import org.jetbrains.kotlin.fir.session.environment.AbstractProjectFileSearchScope
+import org.jetbrains.kotlin.incremental.components.EnumWhenTracker
+import org.jetbrains.kotlin.incremental.components.LookupTracker
+import org.jetbrains.kotlin.js.resolve.JsPlatformAnalyzerServices
+import org.jetbrains.kotlin.library.metadata.resolver.KotlinResolvedLibrary
+import org.jetbrains.kotlin.load.kotlin.PackageAndMetadataPartProvider
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.CommonPlatforms
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.platform.js.JsPlatforms
+import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
+import org.jetbrains.kotlin.platform.konan.NativePlatforms
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
+import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
+import org.jetbrains.kotlin.resolve.konan.platform.NativePlatformAnalyzerServices
+import org.jetbrains.kotlin.resolve.multiplatform.hmppModuleName
+import org.jetbrains.kotlin.resolve.multiplatform.isCommonSource
+
+val isCommonSourceForPsi: (KtFile) -> Boolean = { it.isCommonSource == true }
+val fileBelongsToModuleForPsi: (KtFile, String) -> Boolean = { file, moduleName -> file.hmppModuleName == moduleName }
+
+val GroupedKtSources.isCommonSourceForLt: (KtSourceFile) -> Boolean
+    get() = { it in commonSources }
+
+val GroupedKtSources.fileBelongsToModuleForLt: (KtSourceFile, String) -> Boolean
+    get() = { file, moduleName -> sourcesByModuleName[moduleName].orEmpty().contains(file) }
+
+/**
+ * Creates library session and sources session for JVM platform
+ * Number of created session depends on mode of MPP:
+ *   - disabled
+ *   - legacy (one platform and one common module)
+ *   - HMPP (multiple number of modules)
+ */
+fun <F> prepareJvmSessions(
+    files: List<F>,
+    configuration: CompilerConfiguration,
+    projectEnvironment: AbstractProjectEnvironment,
+    rootModuleName: String,
+    extensionRegistrars: List<FirExtensionRegistrar>,
+    librariesScope: AbstractProjectFileSearchScope,
+    libraryList: DependencyListForCliModule,
+    isCommonSource: (F) -> Boolean,
+    fileBelongsToModule: (F, String) -> Boolean,
+    createProviderAndScopeForIncrementalCompilation: (List<F>) -> IncrementalCompilationContext?,
+): List<SessionWithSources<F>> {
+    val javaSourcesScope = projectEnvironment.getSearchScopeForProjectJavaSources()
+
+    class Context(
+        val javaSourcesScope: AbstractProjectFileSearchScope,
+        val lookupTracker: LookupTracker?,
+        val enumWhenTracker: EnumWhenTracker?
+    )
+
+    return prepareSessions(
+        files, configuration, rootModuleName, JvmPlatforms.unspecifiedJvmPlatform,
+        JvmPlatformAnalyzerServices, libraryList, isCommonSource, fileBelongsToModule,
+        createContext = {
+            val lookupTracker = configuration.get(CommonConfigurationKeys.LOOKUP_TRACKER)
+            val enumWhenTracker = configuration.get(CommonConfigurationKeys.ENUM_WHEN_TRACKER)
+            Context(javaSourcesScope, lookupTracker, enumWhenTracker)
+        },
+        createLibrarySession = { sessionProvider ->
+            FirJvmSessionFactory.createLibrarySession(
+                Name.identifier(rootModuleName),
+                sessionProvider,
+                libraryList.moduleDataProvider,
+                projectEnvironment,
+                librariesScope,
+                projectEnvironment.getPackagePartProvider(librariesScope),
+                configuration.languageVersionSettings,
+                registerExtraComponents = {},
+            )
+
+        }
+    ) { context, moduleFiles, moduleData, sessionProvider, sessionConfigurator ->
+        FirJvmSessionFactory.createModuleBasedSession(
+            moduleData,
+            sessionProvider,
+            context.javaSourcesScope,
+            projectEnvironment,
+            createProviderAndScopeForIncrementalCompilation(moduleFiles),
+            extensionRegistrars,
+            configuration.languageVersionSettings,
+            context.lookupTracker,
+            context.enumWhenTracker,
+            needRegisterJavaElementFinder = true,
+            registerExtraComponents = {},
+            sessionConfigurator,
+        )
+    }
+}
+
+/**
+ * Creates library session and sources session for JS platform
+ * Number of created session depends on mode of MPP:
+ *   - disabled
+ *   - legacy (one platform and one common module)
+ *   - HMPP (multiple number of modules)
+ */
+fun <F> prepareJsSessions(
+    files: List<F>,
+    configuration: CompilerConfiguration,
+    rootModuleName: String,
+    resolvedLibraries: List<KotlinResolvedLibrary>,
+    libraryList: DependencyListForCliModule,
+    extensionRegistrars: List<FirExtensionRegistrar>,
+    isCommonSource: (F) -> Boolean,
+    fileBelongsToModule: (F, String) -> Boolean,
+): List<SessionWithSources<F>> {
+    return prepareSessions(
+        files, configuration, rootModuleName, JsPlatforms.defaultJsPlatform, JsPlatformAnalyzerServices,
+        libraryList, isCommonSource, fileBelongsToModule,
+        createContext = { null },
+        createLibrarySession = { sessionProvider ->
+            FirJsSessionFactory.createLibrarySession(
+                Name.identifier(rootModuleName),
+                resolvedLibraries,
+                sessionProvider,
+                libraryList.moduleDataProvider,
+                configuration.languageVersionSettings,
+                registerExtraComponents = {},
+            )
+        }
+    ) { _, _, moduleData, sessionProvider, sessionConfigurator ->
+        FirJsSessionFactory.createModuleBasedSession(
+            moduleData,
+            sessionProvider,
+            extensionRegistrars,
+            configuration.languageVersionSettings,
+            null,
+            registerExtraComponents = {},
+            init = sessionConfigurator,
+        )
+    }
+}
+
+/**
+ * Creates library session and sources session for Native platform
+ * Number of created session depends on mode of MPP:
+ *   - disabled
+ *   - legacy (one platform and one common module)
+ *   - HMPP (multiple number of modules)
+ */
+fun <F> prepareNativeSessions(
+    files: List<F>,
+    configuration: CompilerConfiguration,
+    rootModuleName: String,
+    resolvedLibraries: List<KotlinResolvedLibrary>,
+    libraryList: DependencyListForCliModule,
+    extensionRegistrars: List<FirExtensionRegistrar>,
+    isCommonSource: (F) -> Boolean,
+    fileBelongsToModule: (F, String) -> Boolean,
+): List<SessionWithSources<F>> {
+    return prepareSessions(
+        files, configuration, rootModuleName, NativePlatforms.unspecifiedNativePlatform, NativePlatformAnalyzerServices,
+        libraryList, isCommonSource, fileBelongsToModule, createContext = { null },
+        createLibrarySession = { sessionProvider ->
+            FirNativeSessionFactory.createLibrarySession(
+                Name.identifier(rootModuleName),
+                resolvedLibraries,
+                sessionProvider,
+                libraryList.moduleDataProvider,
+                configuration.languageVersionSettings,
+                registerExtraComponents = {},
+            )
+        }
+    ) { _, _, moduleData, sessionProvider, sessionConfigurator ->
+        FirNativeSessionFactory.createModuleBasedSession(
+            moduleData,
+            sessionProvider,
+            extensionRegistrars,
+            configuration.languageVersionSettings,
+            sessionConfigurator,
+        )
+    }
+}
+
+/**
+ * Creates library session and sources session for Common platform (for metadata compilation)
+ * Number of created session depends on mode of MPP:
+ *   - disabled
+ *   - legacy (one platform and one common module)
+ *   - HMPP (multiple number of modules)
+ */
+fun <F> prepareCommonSessions(
+    files: List<F>,
+    configuration: CompilerConfiguration,
+    projectEnvironment: AbstractProjectEnvironment,
+    rootModuleName: String,
+    extensionRegistrars: List<FirExtensionRegistrar>,
+    librariesScope: AbstractProjectFileSearchScope,
+    libraryList: DependencyListForCliModule,
+    isCommonSource: (F) -> Boolean,
+    fileBelongsToModule: (F, String) -> Boolean,
+    createProviderAndScopeForIncrementalCompilation: (List<F>) -> IncrementalCompilationContext?,
+): List<SessionWithSources<F>> {
+    return prepareSessions(
+        files, configuration, rootModuleName, CommonPlatforms.defaultCommonPlatform, CommonPlatformAnalyzerServices,
+        libraryList, isCommonSource, fileBelongsToModule, createContext = { null },
+        createLibrarySession = { sessionProvider ->
+            FirCommonSessionFactory.createLibrarySession(
+                Name.identifier(rootModuleName),
+                sessionProvider,
+                libraryList.moduleDataProvider,
+                projectEnvironment,
+                librariesScope,
+                projectEnvironment.getPackagePartProvider(librariesScope) as PackageAndMetadataPartProvider,
+                configuration.languageVersionSettings,
+                registerExtraComponents = {},
+            )
+        }
+    ) { _, moduleFiles, moduleData, sessionProvider, sessionConfigurator ->
+        FirCommonSessionFactory.createModuleBasedSession(
+            moduleData,
+            sessionProvider,
+            librariesScope,
+            projectEnvironment,
+            incrementalCompilationContext = createProviderAndScopeForIncrementalCompilation(moduleFiles),
+            extensionRegistrars,
+            configuration.languageVersionSettings,
+            lookupTracker = configuration.get(CommonConfigurationKeys.LOOKUP_TRACKER),
+            enumWhenTracker = configuration.get(CommonConfigurationKeys.ENUM_WHEN_TRACKER),
+            needRegisterJavaElementFinder = true,
+            registerExtraComponents = {},
+            init = sessionConfigurator
+        )
+    }
+}
+
+// ---------------------------------------------------- Implementation ----------------------------------------------------
+
+private typealias FirSessionProducer<F, C> = (C, List<F>, FirModuleData, FirProjectSessionProvider, FirSessionConfigurator.() -> Unit) -> FirSession
+
+private inline fun <F, C> prepareSessions(
+    files: List<F>,
+    configuration: CompilerConfiguration,
+    rootModuleName: String,
+    targetPlatform: TargetPlatform,
+    analyzerServices: PlatformDependentAnalyzerServices,
+    libraryList: DependencyListForCliModule,
+    isCommonSource: (F) -> Boolean,
+    fileBelongsToModule: (F, String) -> Boolean,
+    createContext: () -> C,
+    createLibrarySession: (FirProjectSessionProvider) -> FirSession,
+    createSourceSession: FirSessionProducer<F, C>,
+): List<SessionWithSources<F>> {
+    val languageVersionSettings = configuration.languageVersionSettings
+
+    val isMppEnabled = languageVersionSettings.supportsFeature(LanguageFeature.MultiPlatformProjects)
+    val hmppModuleStructure = configuration.get(CommonConfigurationKeys.HMPP_MODULE_STRUCTURE)
+    val sessionProvider = FirProjectSessionProvider()
+
+    val context = createContext()
+
+    createLibrarySession(sessionProvider)
+    val extendedAnalysisMode = configuration.getBoolean(CommonConfigurationKeys.USE_FIR_EXTENDED_CHECKERS)
+    val sessionConfigurator: FirSessionConfigurator.() -> Unit = {
+        if (extendedAnalysisMode) {
+            registerExtendedCommonCheckers()
+        }
+    }
+
+    return when {
+        !isMppEnabled -> listOf(
+            createSessionForNonMppProject(
+                context, files, rootModuleName, libraryList, targetPlatform, analyzerServices,
+                sessionProvider, sessionConfigurator, createSourceSession
+            )
+        )
+
+        hmppModuleStructure == null -> createSessionsForLegacyMppProject(
+            context, files, rootModuleName, libraryList, targetPlatform, analyzerServices,
+            sessionProvider, sessionConfigurator, isCommonSource, createSourceSession
+        )
+
+        else -> createSessionsForHmppProject(
+            context, files, hmppModuleStructure, libraryList, targetPlatform, analyzerServices,
+            sessionProvider, sessionConfigurator, fileBelongsToModule, createSourceSession
+        )
+    }
+}
+
+private inline fun <F, C> createSessionForNonMppProject(
+    context: C,
+    files: List<F>,
+    rootModuleName: String,
+    libraryList: DependencyListForCliModule,
+    targetPlatform: TargetPlatform,
+    analyzerServices: PlatformDependentAnalyzerServices,
+    sessionProvider: FirProjectSessionProvider,
+    noinline sessionConfigurator: FirSessionConfigurator.() -> Unit,
+    createFirSession: FirSessionProducer<F, C>,
+): SessionWithSources<F> {
+    val platformModuleData = FirModuleDataImpl(
+        Name.identifier(rootModuleName),
+        libraryList.regularDependencies,
+        dependsOnDependencies = emptyList(),
+        libraryList.friendsDependencies,
+        targetPlatform,
+        analyzerServices
+    )
+
+    val session = createFirSession(context, files, platformModuleData, sessionProvider, sessionConfigurator)
+    return SessionWithSources(session, files)
+}
+
+private inline fun <F, C> createSessionsForLegacyMppProject(
+    context: C,
+    files: List<F>,
+    rootModuleName: String,
+    libraryList: DependencyListForCliModule,
+    targetPlatform: TargetPlatform,
+    analyzerServices: PlatformDependentAnalyzerServices,
+    sessionProvider: FirProjectSessionProvider,
+    noinline sessionConfigurator: FirSessionConfigurator.() -> Unit,
+    isCommonSource: (F) -> Boolean,
+    createFirSession: FirSessionProducer<F, C>,
+): List<SessionWithSources<F>> {
+    val commonModuleData = FirModuleDataImpl(
+        Name.identifier("${rootModuleName}-common"),
+        libraryList.regularDependencies,
+        listOf(),
+        libraryList.friendsDependencies,
+        targetPlatform,
+        analyzerServices
+    )
+
+    val platformModuleData = FirModuleDataImpl(
+        Name.identifier(rootModuleName),
+        libraryList.regularDependencies,
+        listOf(commonModuleData),
+        libraryList.friendsDependencies,
+        targetPlatform,
+        analyzerServices
+    )
+
+    val commonFiles = mutableListOf<F>()
+    val platformFiles = mutableListOf<F>()
+    for (file in files) {
+        (if (isCommonSource(file)) commonFiles else platformFiles).add(file)
+    }
+
+    val commonSession = createFirSession(context, commonFiles, commonModuleData, sessionProvider, sessionConfigurator)
+    val platformSession = createFirSession(context, platformFiles, platformModuleData, sessionProvider, sessionConfigurator)
+
+    return listOf(
+        SessionWithSources(commonSession, commonFiles),
+        SessionWithSources(platformSession, platformFiles)
+    )
+}
+
+private inline fun <F, C> createSessionsForHmppProject(
+    context: C,
+    files: List<F>,
+    hmppModuleStructure: HmppCliModuleStructure,
+    libraryList: DependencyListForCliModule,
+    targetPlatform: TargetPlatform,
+    analyzerServices: PlatformDependentAnalyzerServices,
+    sessionProvider: FirProjectSessionProvider,
+    noinline sessionConfigurator: FirSessionConfigurator.() -> Unit,
+    fileBelongsToModule: (F, String) -> Boolean,
+    createFirSession: FirSessionProducer<F, C>,
+): List<SessionWithSources<F>> {
+    val moduleDataForHmppModule = LinkedHashMap<HmppCliModule, FirModuleData>()
+
+    for (module in hmppModuleStructure.modules) {
+        val dependencies = hmppModuleStructure.dependenciesMap[module]
+            ?.map { moduleDataForHmppModule.getValue(it) }
+            .orEmpty()
+        val moduleData = FirModuleDataImpl(
+            Name.identifier(module.name),
+            libraryList.regularDependencies,
+            dependsOnDependencies = dependencies,
+            libraryList.friendsDependencies,
+            targetPlatform,
+            analyzerServices
+        )
+        moduleDataForHmppModule[module] = moduleData
+    }
+
+    return hmppModuleStructure.modules.map { module ->
+        val moduleData = moduleDataForHmppModule.getValue(module)
+        val sources = files.filter { fileBelongsToModule(it, module.name) }
+        val session = createFirSession(context, sources, moduleData, sessionProvider, sessionConfigurator)
+        SessionWithSources(session, sources)
+    }
+}
+
+data class SessionWithSources<F>(val session: FirSession, val files: List<F>)
@@ -11,10 +11,7 @@ import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.jvm.JvmGeneratorExtensions
 import org.jetbrains.kotlin.backend.jvm.JvmIrCodegenFactory
 import org.jetbrains.kotlin.backend.jvm.JvmIrDeserializerImpl
-import org.jetbrains.kotlin.cli.common.CLICompiler
-import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
-import org.jetbrains.kotlin.cli.common.CommonCompilerPerformanceManager
-import org.jetbrains.kotlin.cli.common.checkKotlinPackageUsage
+import org.jetbrains.kotlin.cli.common.*
 import org.jetbrains.kotlin.cli.common.fir.FirDiagnosticsCompilerResultsReporter
 import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
@@ -30,15 +27,12 @@ import org.jetbrains.kotlin.diagnostics.impl.BaseDiagnosticsCollector
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.backend.jvm.*
-import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.pipeline.*
 import org.jetbrains.kotlin.fir.session.*
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
-import org.jetbrains.kotlin.fir.session.environment.AbstractProjectFileSearchScope
 import org.jetbrains.kotlin.fir.types.arrayElementType
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.isArrayType
@@ -50,13 +44,12 @@ import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompil
 import org.jetbrains.kotlin.modules.Module
 import org.jetbrains.kotlin.modules.TargetId
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
+import org.jetbrains.kotlin.resolve.multiplatform.hmppModuleName
 import org.jetbrains.kotlin.resolve.multiplatform.isCommonSource
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
+import org.jetbrains.kotlin.utils.addToStdlib.runUnless
 import java.io.File

 object FirKotlinToJvmBytecodeCompiler {
@@ -66,14 +59,15 @@ object FirKotlinToJvmBytecodeCompiler {
        messageCollector: MessageCollector,
        allSources: List<KtFile>,
        buildFile: File?,
-        chunk: List<Module>,
-        extendedAnalysisMode: Boolean
+        chunk: List<Module>
    ): Boolean {
        val performanceManager = projectConfiguration.get(CLIConfigurationKeys.PERF_MANAGER)

        val notSupportedPlugins = mutableListOf<String?>().apply {
-            projectConfiguration.get(ComponentRegistrar.PLUGIN_COMPONENT_REGISTRARS).collectIncompatiblePluginNamesTo(this, ComponentRegistrar::supportsK2)
-            projectConfiguration.get(CompilerPluginRegistrar.COMPILER_PLUGIN_REGISTRARS).collectIncompatiblePluginNamesTo(this, CompilerPluginRegistrar::supportsK2)
+            projectConfiguration.get(ComponentRegistrar.PLUGIN_COMPONENT_REGISTRARS)
+                .collectIncompatiblePluginNamesTo(this, ComponentRegistrar::supportsK2)
+            projectConfiguration.get(CompilerPluginRegistrar.COMPILER_PLUGIN_REGISTRARS)
+                .collectIncompatiblePluginNamesTo(this, CompilerPluginRegistrar::supportsK2)
        }

        if (notSupportedPlugins.isNotEmpty()) {
@@ -109,8 +103,7 @@ object FirKotlinToJvmBytecodeCompiler {
                performanceManager,
                targetIds,
                incrementalComponents,
-                extendedAnalysisMode,
-                firExtensionRegistrars = project?.let { FirExtensionRegistrar.getInstances(it) } ?: emptyList(),
+                extensionRegistrars = project?.let { FirExtensionRegistrar.getInstances(it) } ?: emptyList(),
                irGenerationExtensions = project?.let { IrGenerationExtension.getInstances(it) } ?: emptyList()
            )
            val generationState = context.compileModule() ?: return false
@@ -118,7 +111,7 @@ object FirKotlinToJvmBytecodeCompiler {
        }

        val mainClassFqName: FqName? = runIf(chunk.size == 1 && projectConfiguration.get(JVMConfigurationKeys.OUTPUT_JAR) != null) {
-            findMainClass(outputs.single().first.platformOutput.fir)
+            findMainClass(outputs.single().first.outputs.last().fir)
        }

        return writeOutputsIfNeeded(
@@ -204,97 +197,21 @@ object FirKotlinToJvmBytecodeCompiler {
        providerAndScopeForIncrementalCompilation?.precompiledBinariesFileScope?.let {
            librariesScope -= it
        }
-
-        val languageVersionSettings = moduleConfiguration.languageVersionSettings
-
-        val isMppEnabled = languageVersionSettings.supportsFeature(LanguageFeature.MultiPlatformProjects)
-
-        val sessionProvider = FirProjectSessionProvider()
-
-        val moduleName = module.getModuleName()
-        val libraryList = createFirLibraryListAndSession(
-            moduleName, moduleConfiguration, projectEnvironment,
-            scope = librariesScope, librariesScope = librariesScope, friendPaths = module.getFriendPaths(), sessionProvider
+        val rootModuleName = module.getModuleName()
+        val libraryList = createLibraryListForJvm(rootModuleName, moduleConfiguration, module.getFriendPaths())
+        val sessionsWithSources = prepareJvmSessions(
+            ktFiles, moduleConfiguration, projectEnvironment, rootModuleName,
+            extensionRegistrars, librariesScope, libraryList,
+            isCommonSource = { it.isCommonSource == true },
+            fileBelongsToModule = { file, moduleName -> file.hmppModuleName == moduleName },
+            createProviderAndScopeForIncrementalCompilation = { providerAndScopeForIncrementalCompilation }
        )

-        val commonModuleData = runIf(isMppEnabled) {
-            FirModuleDataImpl(
-                Name.identifier("${module.getModuleName()}-common"),
-                libraryList.regularDependencies,
-                listOf(),
-                libraryList.friendsDependencies,
-                JvmPlatforms.unspecifiedJvmPlatform,
-                JvmPlatformAnalyzerServices
-            )
-        }
-        val platformModuleData = FirModuleDataImpl(
-            Name.identifier(module.getModuleName()),
-            libraryList.regularDependencies,
-            listOfNotNull(commonModuleData),
-            libraryList.friendsDependencies,
-            JvmPlatforms.unspecifiedJvmPlatform,
-            JvmPlatformAnalyzerServices
-        )
-
-        val lookupTracker = moduleConfiguration.get(CommonConfigurationKeys.LOOKUP_TRACKER)
-        val enumWhenTracker = moduleConfiguration.get(CommonConfigurationKeys.ENUM_WHEN_TRACKER)
-        val sessionConfigurator: FirSessionConfigurator.() -> Unit = {
-            if (extendedAnalysisMode) {
-                registerExtendedCommonCheckers()
-            }
+        val outputs = sessionsWithSources.map { (session, sources) ->
+            buildResolveAndCheckFir(session, sources, diagnosticsReporter)
        }

-        val commonKtFiles = mutableListOf<KtFile>()
-        val platformKtFiles = mutableListOf<KtFile>()
-        val commonSourcesScope: AbstractProjectFileSearchScope?
-        val platformSourcesScope: AbstractProjectFileSearchScope
-        if (isMppEnabled) {
-            for (ktFile in ktFiles) {
-                (if (ktFile.isCommonSource == true) commonKtFiles else platformKtFiles).add(ktFile)
-            }
-            commonSourcesScope = projectEnvironment.getSearchScopeByPsiFiles(commonKtFiles)
-            platformSourcesScope = sourceScope - commonSourcesScope
-        } else {
-            platformKtFiles.addAll(ktFiles)
-            commonSourcesScope = null
-            platformSourcesScope = sourceScope
-        }
-
-        val commonSession = runIf(isMppEnabled) {
-            FirJvmSessionFactory.createModuleBasedSession(
-                commonModuleData!!,
-                sessionProvider,
-                commonSourcesScope!!,
-                projectEnvironment,
-                providerAndScopeForIncrementalCompilation,
-                firExtensionRegistrars,
-                languageVersionSettings,
-                lookupTracker,
-                enumWhenTracker,
-                needRegisterJavaElementFinder = true,
-                registerExtraComponents = {},
-                sessionConfigurator
-            )
-        }
-        val platformSession = FirJvmSessionFactory.createModuleBasedSession(
-            platformModuleData,
-            sessionProvider,
-            platformSourcesScope,
-            projectEnvironment,
-            providerAndScopeForIncrementalCompilation,
-            firExtensionRegistrars,
-            languageVersionSettings,
-            lookupTracker,
-            enumWhenTracker,
-            needRegisterJavaElementFinder = true,
-            registerExtraComponents = {},
-            sessionConfigurator,
-        )
-
-        val commonOutput = commonSession?.let { buildResolveAndCheckFir(it, commonKtFiles, diagnosticsReporter) }
-        val platformOutput = buildResolveAndCheckFir(platformSession, platformKtFiles, diagnosticsReporter)
-
-        return if (syntaxErrors || diagnosticsReporter.hasErrors) null else FirResult(platformOutput, commonOutput)
+        return runUnless(syntaxErrors || diagnosticsReporter.hasErrors) { FirResult(outputs) }
    }

    private fun CompilationContext.runBackend(
@@ -365,8 +282,7 @@ object FirKotlinToJvmBytecodeCompiler {
        val performanceManager: CommonCompilerPerformanceManager?,
        val targetIds: List<TargetId>?,
        val incrementalComponents: IncrementalCompilationComponents?,
-        val extendedAnalysisMode: Boolean,
-        val firExtensionRegistrars: List<FirExtensionRegistrar>,
+        val extensionRegistrars: List<FirExtensionRegistrar>,
        val irGenerationExtensions: Collection<IrGenerationExtension>
    )
 }
@@ -70,7 +70,6 @@ object KotlinToJVMBytecodeCompiler {

        val projectConfiguration = environment.configuration
        if (projectConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR)) {
-            val extendedAnalysisMode = projectConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR_EXTENDED_CHECKERS)
            val projectEnvironment =
                VfsBasedProjectEnvironment(
                    environment.project,
@@ -81,7 +80,7 @@ object KotlinToJVMBytecodeCompiler {
                environment.configuration,
                environment.messageCollector,
                environment.getSourceFiles(),
-                buildFile, chunk, extendedAnalysisMode
+                buildFile, chunk
            )
        }

@@ -10,7 +10,6 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileSystem
-import org.jetbrains.kotlin.analyzer.common.CommonPlatformAnalyzerServices
 import org.jetbrains.kotlin.backend.common.output.OutputFileCollection
 import org.jetbrains.kotlin.backend.common.output.SimpleOutputFileCollection
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
@@ -27,17 +26,12 @@ import org.jetbrains.kotlin.codegen.state.GenerationStateEventCallback
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
-import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.fir.BinaryModuleData
 import org.jetbrains.kotlin.fir.DependencyListForCliModule
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
-import org.jetbrains.kotlin.fir.session.FirCommonSessionFactory
-import org.jetbrains.kotlin.fir.session.FirJvmSessionFactory
 import org.jetbrains.kotlin.fir.session.IncrementalCompilationContext
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectFileSearchScope
 import org.jetbrains.kotlin.javac.JavacWrapper
-import org.jetbrains.kotlin.load.kotlin.PackageAndMetadataPartProvider
 import org.jetbrains.kotlin.load.kotlin.incremental.IncrementalPackagePartProvider
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
 import org.jetbrains.kotlin.modules.JavaRootPath
@@ -45,7 +39,6 @@ import org.jetbrains.kotlin.modules.Module
 import org.jetbrains.kotlin.modules.TargetId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.platform.CommonPlatforms
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.psi.KtFile
@@ -226,20 +219,15 @@ fun createContextForIncrementalCompilation(
    return IncrementalCompilationContext(emptyList(), packagePartProvider, incrementalCompilationScope)
 }

-fun createFirLibraryListAndSession(
+fun createLibraryListForJvm(
    moduleName: String,
    configuration: CompilerConfiguration,
-    projectEnvironment: AbstractProjectEnvironment,
-    scope: AbstractProjectFileSearchScope,
-    librariesScope: AbstractProjectFileSearchScope,
-    friendPaths: List<String>,
-    sessionProvider: FirProjectSessionProvider,
-    isJvm: Boolean = true
+    friendPaths: List<String>
 ): DependencyListForCliModule {
    val binaryModuleData = BinaryModuleData.initialize(
        Name.identifier(moduleName),
-        if (isJvm) JvmPlatforms.unspecifiedJvmPlatform else CommonPlatforms.defaultCommonPlatform,
-        if (isJvm) JvmPlatformAnalyzerServices else CommonPlatformAnalyzerServices
+        JvmPlatforms.unspecifiedJvmPlatform,
+        JvmPlatformAnalyzerServices
    )
    val libraryList = DependencyListForCliModule.build(binaryModuleData) {
        dependencies(configuration.jvmClasspathRoots.map { it.toPath() })
@@ -247,29 +235,6 @@ fun createFirLibraryListAndSession(
        friendDependencies(configuration[JVMConfigurationKeys.FRIEND_PATHS] ?: emptyList())
        friendDependencies(friendPaths)
    }
-    if (isJvm) {
-        FirJvmSessionFactory.createLibrarySession(
-            Name.identifier(moduleName),
-            sessionProvider,
-            libraryList.moduleDataProvider,
-            projectEnvironment,
-            scope,
-            projectEnvironment.getPackagePartProvider(librariesScope),
-            configuration.languageVersionSettings,
-            registerExtraComponents = {},
-        )
-    } else {
-        FirCommonSessionFactory.createLibrarySession(
-            Name.identifier(moduleName),
-            sessionProvider,
-            libraryList.moduleDataProvider,
-            projectEnvironment,
-            scope,
-            projectEnvironment.getPackagePartProvider(librariesScope) as PackageAndMetadataPartProvider,
-            configuration.languageVersionSettings,
-            registerExtraComponents = {},
-        )
-    }
    return libraryList
 }

@@ -22,8 +22,7 @@ import org.jetbrains.kotlin.KtVirtualFileSourceFile
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.jvm.JvmIrCodegenFactory
 import org.jetbrains.kotlin.backend.jvm.JvmIrDeserializerImpl
-import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
-import org.jetbrains.kotlin.cli.common.CommonCompilerPerformanceManager
+import org.jetbrains.kotlin.cli.common.*
 import org.jetbrains.kotlin.cli.common.config.kotlinSourceRoots
 import org.jetbrains.kotlin.cli.common.fir.reportToMessageCollector
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
@@ -40,21 +39,19 @@ import org.jetbrains.kotlin.cli.jvm.modules.CliJavaModuleResolver
 import org.jetbrains.kotlin.codegen.ClassBuilderFactories
 import org.jetbrains.kotlin.codegen.CodegenFactory
 import org.jetbrains.kotlin.codegen.state.GenerationState
-import org.jetbrains.kotlin.config.*
+import org.jetbrains.kotlin.config.CommonConfigurationKeys
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
-import org.jetbrains.kotlin.fir.FirModuleDataImpl
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmBackendClassResolver
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmBackendExtension
 import org.jetbrains.kotlin.fir.backend.jvm.JvmFir2IrExtensions
-import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.pipeline.*
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.session.FirJvmSessionFactory
-import org.jetbrains.kotlin.fir.session.FirSessionConfigurator
 import org.jetbrains.kotlin.fir.session.IncrementalCompilationContext
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectEnvironment
 import org.jetbrains.kotlin.fir.session.environment.AbstractProjectFileSearchScope
@@ -68,14 +65,11 @@ import org.jetbrains.kotlin.load.kotlin.incremental.IncrementalPackagePartProvid
 import org.jetbrains.kotlin.modules.Module
 import org.jetbrains.kotlin.modules.TargetId
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.CommonPlatforms
 import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.resolve.ModuleAnnotationsResolver
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
-import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
-import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import java.io.File
 import kotlin.reflect.KFunction2

@@ -104,13 +98,13 @@ fun compileModulesUsingFrontendIrAndLightTree(
        val moduleConfiguration = compilerConfiguration.copy().applyModuleProperties(module, buildFile).apply {
            put(JVMConfigurationKeys.FRIEND_PATHS, module.getFriendPaths())
        }
-        val (platformSources, commonSources, sourcesByModuleName) = collectSources(compilerConfiguration, projectEnvironment, messageCollector)
+        val groupedSources = collectSources(compilerConfiguration, projectEnvironment, messageCollector)

        val compilerInput = ModuleCompilerInput(
            TargetId(module),
-            CommonPlatforms.defaultCommonPlatform, commonSources,
-            JvmPlatforms.unspecifiedJvmPlatform, platformSources,
-            sourcesByModuleName,
+            groupedSources,
+            CommonPlatforms.defaultCommonPlatform,
+            JvmPlatforms.unspecifiedJvmPlatform,
            moduleConfiguration
        )

@@ -133,7 +127,7 @@ fun compileModulesUsingFrontendIrAndLightTree(

        // TODO: consider what to do if many modules has main classes
        if (mainClassFqName == null && moduleConfiguration.get(JVMConfigurationKeys.OUTPUT_JAR) != null) {
-            mainClassFqName = findMainClass(analysisResults.platformOutput.fir)
+            mainClassFqName = findMainClass(analysisResults.outputs.last().fir)
        }

        if (diagnosticsReporter.hasErrors) {
@@ -171,8 +165,8 @@ fun compileModulesUsingFrontendIrAndLightTree(
 }

 data class GroupedKtSources(
-    val platformSources: Set<KtSourceFile>,
-    val commonSources: Set<KtSourceFile>,
+    val platformSources: Collection<KtSourceFile>,
+    val commonSources: Collection<KtSourceFile>,
    val sourcesByModuleName: Map<String, Set<KtSourceFile>>,
 )

@@ -187,7 +181,8 @@ fun collectSources(

    // TODO: the scripts checking should be part of the scripting plugin functionality, as it is implemented now in ScriptingProcessSourcesBeforeCompilingExtension
    // TODO: implement in the next round of K2 scripting support (https://youtrack.jetbrains.com/issue/KT-55728)
-    val skipScriptsInLtMode = compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR) && compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE)
+    val skipScriptsInLtMode = compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR) &&
+            compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE)
    var skipScriptsInLtModeWarning = false

    compilerConfiguration.kotlinSourceRoots.forAllFiles(
@@ -308,122 +303,48 @@ fun compileModuleToAnalyzedFir(
    diagnosticsReporter: DiagnosticReporter,
    performanceManager: CommonCompilerPerformanceManager?
 ): FirResult {
-    val languageVersionSettings = input.configuration.languageVersionSettings
    val projectEnvironment = environment.projectEnvironment
    val moduleConfiguration = input.configuration
-    val isMppEnabled = languageVersionSettings.supportsFeature(LanguageFeature.MultiPlatformProjects)
-    val sourcesScope = environment.projectEnvironment.getSearchScopeBySourceFiles(input.platformSources)
-
-    val commonSourcesScope: AbstractProjectFileSearchScope?
-    val platformSourcesScope: AbstractProjectFileSearchScope
-    if (isMppEnabled) {
-        commonSourcesScope = projectEnvironment.getSearchScopeBySourceFiles(input.commonSources)
-        platformSourcesScope = sourcesScope - commonSourcesScope
-    } else {
-        commonSourcesScope = null
-        platformSourcesScope = sourcesScope
-    }

    var librariesScope = projectEnvironment.getSearchScopeForProjectLibraries()
-    val commonProviderAndScopeForIncrementalCompilation = runIf(isMppEnabled) {
-        createContextForIncrementalCompilation(
-            moduleConfiguration,
-            projectEnvironment,
-            commonSourcesScope!!,
-            previousStepsSymbolProviders,
-            incrementalExcludesScope
-        )?.also { (_, _, precompiledBinariesFileScope) ->
-            precompiledBinariesFileScope?.let { librariesScope -= it }
-        }
-    }
-    val platformProviderAndScopeForIncrementalCompilation = createContextForIncrementalCompilation(
-        moduleConfiguration,
-        projectEnvironment,
-        platformSourcesScope,
-        previousStepsSymbolProviders,
-        incrementalExcludesScope
-    )?.also { (_, _, precompiledBinariesFileScope) ->
-        precompiledBinariesFileScope?.let { librariesScope -= it }
-    }
+    val rootModuleName = input.targetId.name

-    val sessionProvider = FirProjectSessionProvider()
-    val moduleName = input.targetId.name
-
-    val libraryList = createFirLibraryListAndSession(
-        moduleName, input.configuration, projectEnvironment,
-        scope = projectEnvironment.getSearchScopeForProjectLibraries(),
-        librariesScope = librariesScope,
-        friendPaths = emptyList(),
-        sessionProvider
-    )
-
-    val commonModuleData = runIf(isMppEnabled) {
-        FirModuleDataImpl(
-            Name.identifier("${moduleName}-common"),
-            libraryList.regularDependencies,
-            listOf(),
-            libraryList.friendsDependencies,
-            JvmPlatforms.unspecifiedJvmPlatform,
-            JvmPlatformAnalyzerServices
-        )
-    }
-    val platformModuleData = FirModuleDataImpl(
-        Name.identifier(moduleName),
-        libraryList.regularDependencies,
-        listOfNotNull(commonModuleData),
-        libraryList.friendsDependencies,
-        JvmPlatforms.unspecifiedJvmPlatform,
-        JvmPlatformAnalyzerServices
-    )
-
-    val extensionRegistrars = (projectEnvironment as? VfsBasedProjectEnvironment)?.let { FirExtensionRegistrar.getInstances(it.project) }
+    val extensionRegistrars = (projectEnvironment as? VfsBasedProjectEnvironment)
+        ?.let { FirExtensionRegistrar.getInstances(it.project) }
        ?: emptyList()
-    val lookupTracker = moduleConfiguration.get(CommonConfigurationKeys.LOOKUP_TRACKER)
-    val enumWhenTracker = moduleConfiguration.get(CommonConfigurationKeys.ENUM_WHEN_TRACKER)
-    val extendedAnalysisMode = input.configuration.getBoolean(CommonConfigurationKeys.USE_FIR_EXTENDED_CHECKERS)
-    val sessionConfigurator: FirSessionConfigurator.() -> Unit = {
-        if (extendedAnalysisMode) {
-            registerExtendedCommonCheckers()
-        }
-    }
-    val javaSourcesScope = projectEnvironment.getSearchScopeForProjectJavaSources()

-    val commonSession = runIf(isMppEnabled) {
-        FirJvmSessionFactory.createModuleBasedSession(
-            commonModuleData!!,
-            sessionProvider,
-            javaSourcesScope,
-            projectEnvironment,
-            commonProviderAndScopeForIncrementalCompilation,
-            extensionRegistrars,
-            languageVersionSettings,
-            lookupTracker,
-            enumWhenTracker,
-            needRegisterJavaElementFinder = true,
-            registerExtraComponents = {},
-            sessionConfigurator,
-        )
+    val allSources = mutableListOf<KtSourceFile>().apply {
+        addAll(input.groupedSources.commonSources)
+        addAll(input.groupedSources.platformSources)
    }
-    val platformSession = FirJvmSessionFactory.createModuleBasedSession(
-        platformModuleData,
-        sessionProvider,
-        javaSourcesScope,
-        projectEnvironment,
-        platformProviderAndScopeForIncrementalCompilation,
-        extensionRegistrars,
-        languageVersionSettings,
-        lookupTracker,
-        enumWhenTracker,
-        needRegisterJavaElementFinder = true,
-        registerExtraComponents = {},
-        sessionConfigurator
+    // TODO: handle friends paths
+    val libraryList = createLibraryListForJvm(rootModuleName, moduleConfiguration, friendPaths = emptyList())
+    val sessionWithSources = prepareJvmSessions(
+        allSources, moduleConfiguration, projectEnvironment, rootModuleName,
+        extensionRegistrars, librariesScope, libraryList,
+        isCommonSource = input.groupedSources.isCommonSourceForLt,
+        fileBelongsToModule = input.groupedSources.fileBelongsToModuleForLt,
+        createProviderAndScopeForIncrementalCompilation = { files ->
+            val scope = projectEnvironment.getSearchScopeBySourceFiles(files)
+            createContextForIncrementalCompilation(
+                moduleConfiguration,
+                projectEnvironment,
+                scope,
+                previousStepsSymbolProviders,
+                incrementalExcludesScope
+            )?.also { (_, _, precompiledBinariesFileScope) ->
+                precompiledBinariesFileScope?.let { librariesScope -= it }
+            }
+        }
    )

    val countFilesAndLines = if (performanceManager == null) null else performanceManager::addSourcesStats
-    val commonOutput = commonSession?.let { buildResolveAndCheckFir(it, input.commonSources, diagnosticsReporter, countFilesAndLines) }
-    val platformOutput = buildResolveAndCheckFir(platformSession, input.platformSources, diagnosticsReporter, countFilesAndLines)

-    return FirResult(platformOutput, commonOutput)
+    val outputs = sessionWithSources.map { (session, sources) ->
+        buildResolveAndCheckFir(session, sources, diagnosticsReporter, countFilesAndLines)
+    }
+
+    return FirResult(outputs)
 }

 private fun buildResolveAndCheckFir(
@@ -5,7 +5,6 @@

 package org.jetbrains.kotlin.cli.jvm.compiler.pipeline

-import org.jetbrains.kotlin.KtSourceFile
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.diagnostics.impl.BaseDiagnosticsCollector
@@ -22,11 +21,9 @@ import org.jetbrains.kotlin.platform.TargetPlatform

 data class ModuleCompilerInput(
    val targetId: TargetId,
+    val groupedSources: GroupedKtSources,
    val commonPlatform: TargetPlatform,
-    val commonSources: Collection<KtSourceFile>,
    val platform: TargetPlatform,
-    val platformSources: Collection<KtSourceFile>,
-    val sourcesByModule: Map<String, Collection<KtSourceFile>>,
    val configuration: CompilerConfiguration,
    val friendFirModules: Collection<FirModuleData> = emptyList()
 )
@@ -7,22 +7,25 @@ package org.jetbrains.kotlin.cli.metadata

 import com.intellij.openapi.vfs.StandardFileSystems
 import com.intellij.openapi.vfs.VirtualFileManager
-import org.jetbrains.kotlin.KtSourceFile
 import org.jetbrains.kotlin.analyzer.common.CommonPlatformAnalyzerServices
-import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
+import org.jetbrains.kotlin.cli.common.*
 import org.jetbrains.kotlin.cli.common.fir.FirDiagnosticsCompilerResultsReporter
-import org.jetbrains.kotlin.cli.jvm.compiler.*
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.VfsBasedProjectEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.createContextForIncrementalCompilation
 import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.collectSources
 import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.createContextForIncrementalCompilation
+import org.jetbrains.kotlin.cli.jvm.compiler.toAbstractProjectEnvironment
+import org.jetbrains.kotlin.cli.jvm.config.jvmClasspathRoots
+import org.jetbrains.kotlin.cli.jvm.config.jvmModularRoots
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
-import org.jetbrains.kotlin.fir.FirModuleDataImpl
-import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
+import org.jetbrains.kotlin.fir.BinaryModuleData
+import org.jetbrains.kotlin.fir.DependencyListForCliModule
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
-import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.pipeline.ModuleCompilerAnalyzedOutput
@@ -32,23 +35,19 @@ import org.jetbrains.kotlin.fir.pipeline.resolveAndCheckFir
 import org.jetbrains.kotlin.fir.serialization.FirElementAwareSerializableStringTable
 import org.jetbrains.kotlin.fir.serialization.FirKLibSerializerExtension
 import org.jetbrains.kotlin.fir.serialization.serializeSingleFirFile
-import org.jetbrains.kotlin.fir.session.FirCommonSessionFactory
-import org.jetbrains.kotlin.fir.session.IncrementalCompilationContext
-import org.jetbrains.kotlin.fir.session.environment.AbstractProjectFileSearchScope
 import org.jetbrains.kotlin.library.SerializedMetadata
 import org.jetbrains.kotlin.library.metadata.KlibMetadataHeaderFlags
 import org.jetbrains.kotlin.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.modules.TargetId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.CommonPlatforms
-import org.jetbrains.kotlin.psi.KtFile
 import java.io.File

 internal class FirMetadataSerializer(
    configuration: CompilerConfiguration,
    environment: KotlinCoreEnvironment
-) : AbstractMetadataSerializer<ModuleCompilerAnalyzedOutput>(configuration, environment) {
-    override fun analyze(): ModuleCompilerAnalyzedOutput? {
+) : AbstractMetadataSerializer<List<ModuleCompilerAnalyzedOutput>>(configuration, environment) {
+    override fun analyze(): List<ModuleCompilerAnalyzedOutput>? {
        val performanceManager = environment.configuration.getNotNull(CLIConfigurationKeys.PERF_MANAGER)
        performanceManager.notifyAnalysisStarted()

@@ -57,23 +56,28 @@ internal class FirMetadataSerializer(
        val moduleName = Name.special("<${configuration.getNotNull(CommonConfigurationKeys.MODULE_NAME)}>")
        val isLightTree = configuration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE)

-        val sessionProvider = FirProjectSessionProvider()
+        val rootModuleName = moduleName.asString()
+        val binaryModuleData = BinaryModuleData.initialize(
+            Name.identifier(rootModuleName),
+            CommonPlatforms.defaultCommonPlatform,
+            CommonPlatformAnalyzerServices
+        )
+        val libraryList = DependencyListForCliModule.build(binaryModuleData) {
+            dependencies(configuration.jvmClasspathRoots.map { it.toPath() })
+            dependencies(configuration.jvmModularRoots.map { it.toPath() })
+            friendDependencies(configuration[JVMConfigurationKeys.FRIEND_PATHS] ?: emptyList())
+        }

-        val projectEnvironment: VfsBasedProjectEnvironment
-        var librariesScope: AbstractProjectFileSearchScope
-        val sourceScope: AbstractProjectFileSearchScope
-        val librariesHelperScope: AbstractProjectFileSearchScope
-        var psiFiles: List<KtFile>? = null
-        var ltFiles: List<KtSourceFile>? = null
-        val providerAndScopeForIncrementalCompilation: IncrementalCompilationContext?
+        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()

-        if (isLightTree) {
-            projectEnvironment = environment.toAbstractProjectEnvironment() as VfsBasedProjectEnvironment
-            librariesScope = projectEnvironment.getSearchScopeForProjectLibraries()
-            librariesHelperScope = projectEnvironment.getSearchScopeForProjectLibraries()
-            ltFiles = collectSources(configuration, projectEnvironment, messageCollector).let { it.commonSources + it.platformSources }.toList()
-            sourceScope = projectEnvironment.getSearchScopeBySourceFiles(ltFiles)
-            providerAndScopeForIncrementalCompilation = createContextForIncrementalCompilation(
+        val outputs = if (isLightTree) {
+            val projectEnvironment = environment.toAbstractProjectEnvironment() as VfsBasedProjectEnvironment
+            var librariesScope = projectEnvironment.getSearchScopeForProjectLibraries()
+            val groupedSources = collectSources(configuration, projectEnvironment, messageCollector)
+            val extensionRegistrars = FirExtensionRegistrar.getInstances(projectEnvironment.project)
+            val ltFiles = groupedSources.let { it.commonSources + it.platformSources }.toList()
+            val sourceScope = projectEnvironment.getSearchScopeBySourceFiles(ltFiles)
+            val providerAndScopeForIncrementalCompilation = createContextForIncrementalCompilation(
                configuration,
                projectEnvironment,
                sourceScope,
@@ -82,16 +86,26 @@ internal class FirMetadataSerializer(
            )?.also { (_, _, precompiledBinariesFileScope) ->
                precompiledBinariesFileScope?.let { librariesScope -= it }
            }
+            val sessionsWithSources = prepareCommonSessions(
+                ltFiles, configuration, projectEnvironment, rootModuleName, extensionRegistrars,
+                librariesScope, libraryList, groupedSources.isCommonSourceForLt, groupedSources.fileBelongsToModuleForLt,
+                createProviderAndScopeForIncrementalCompilation = { providerAndScopeForIncrementalCompilation }
+            )
+            sessionsWithSources.map { (session, files) ->
+                val firFiles = session.buildFirViaLightTree(files, diagnosticsReporter, performanceManager::addSourcesStats)
+                resolveAndCheckFir(session, firFiles, diagnosticsReporter)
+            }
        } else {
-            projectEnvironment = VfsBasedProjectEnvironment(
+            val projectEnvironment = VfsBasedProjectEnvironment(
                environment.project,
                VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
            ) { environment.createPackagePartProvider(it) }
-            librariesScope = projectEnvironment.getSearchScopeForProjectLibraries()
-            librariesHelperScope = librariesScope
-            psiFiles = environment.getSourceFiles()
-            sourceScope = projectEnvironment.getSearchScopeByPsiFiles(psiFiles) + projectEnvironment.getSearchScopeForProjectJavaSources()
-            providerAndScopeForIncrementalCompilation = createContextForIncrementalCompilation(
+            var librariesScope = projectEnvironment.getSearchScopeForProjectLibraries()
+            val extensionRegistrars = FirExtensionRegistrar.getInstances(projectEnvironment.project)
+            val psiFiles = environment.getSourceFiles()
+            val sourceScope =
+                projectEnvironment.getSearchScopeByPsiFiles(psiFiles) + projectEnvironment.getSearchScopeForProjectJavaSources()
+            val providerAndScopeForIncrementalCompilation = createContextForIncrementalCompilation(
                projectEnvironment,
                configuration.get(JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS),
                configuration,
@@ -101,80 +115,51 @@ internal class FirMetadataSerializer(
            providerAndScopeForIncrementalCompilation?.precompiledBinariesFileScope?.let {
                librariesScope -= it
            }
+            val sessionsWithSources = prepareCommonSessions(
+                psiFiles, configuration, projectEnvironment, rootModuleName, extensionRegistrars,
+                librariesScope, libraryList, isCommonSourceForPsi, fileBelongsToModuleForPsi,
+                createProviderAndScopeForIncrementalCompilation = { providerAndScopeForIncrementalCompilation }
+            )
+
+            sessionsWithSources.map { (session, files) ->
+                val firFiles = session.buildFirFromKtFiles(files)
+                resolveAndCheckFir(session, firFiles, diagnosticsReporter)
+            }
        }

-        val libraryList = createFirLibraryListAndSession(
-            moduleName.asString(), configuration, projectEnvironment,
-            scope = librariesHelperScope, librariesScope = librariesScope, friendPaths = emptyList(), sessionProvider = sessionProvider,
-            isJvm = false
-        )
-
-        val commonModuleData = FirModuleDataImpl(
-            moduleName,
-            libraryList.regularDependencies,
-            listOf(),
-            libraryList.friendsDependencies,
-            CommonPlatforms.defaultCommonPlatform,
-            CommonPlatformAnalyzerServices
-        )
-        val project = projectEnvironment.project
-        val session = FirCommonSessionFactory.createModuleBasedSession(
-            commonModuleData,
-            sessionProvider,
-            sourceScope,
-            projectEnvironment,
-            incrementalCompilationContext = providerAndScopeForIncrementalCompilation,
-            FirExtensionRegistrar.getInstances(project),
-            configuration.languageVersionSettings,
-            lookupTracker = configuration.get(CommonConfigurationKeys.LOOKUP_TRACKER),
-            enumWhenTracker = configuration.get(CommonConfigurationKeys.ENUM_WHEN_TRACKER),
-            needRegisterJavaElementFinder = true,
-            registerExtraComponents = {},
-            init = {
-                if (configuration.getBoolean(CommonConfigurationKeys.USE_FIR_EXTENDED_CHECKERS)) {
-                    registerExtendedCommonCheckers()
-                }
-            }
-        )
-
-        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
-        val firFiles = if (isLightTree)
-            session.buildFirViaLightTree(ltFiles!!, diagnosticsReporter, performanceManager::addSourcesStats)
-        else
-            session.buildFirFromKtFiles(psiFiles!!)
-
-        val result = resolveAndCheckFir(session, firFiles, diagnosticsReporter)

        return if (diagnosticsReporter.hasErrors) {
            val renderDiagnosticNames = configuration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
            FirDiagnosticsCompilerResultsReporter.reportToMessageCollector(diagnosticsReporter, messageCollector, renderDiagnosticNames)
            null
        } else {
-            result
+            outputs
        }.also {
            performanceManager.notifyAnalysisFinished()
        }
    }

-    override fun serialize(analysisResult: ModuleCompilerAnalyzedOutput, destDir: File) {
+    override fun serialize(analysisResult: List<ModuleCompilerAnalyzedOutput>, destDir: File) {
        val fragments = mutableMapOf<String, MutableList<ByteArray>>()

-        val (session, scopeSession, fir) = analysisResult
+        for (output in analysisResult) {
+            val (session, scopeSession, fir) = output

-        val languageVersionSettings = environment.configuration.languageVersionSettings
-        for (firFile in fir) {
-            val packageFragment = serializeSingleFirFile(
-                firFile,
-                session,
-                scopeSession,
-                FirKLibSerializerExtension(session, metadataVersion, FirElementAwareSerializableStringTable()),
-                languageVersionSettings
-            )
-            fragments.getOrPut(firFile.packageFqName.asString()) { mutableListOf() }.add(packageFragment.toByteArray())
+            val languageVersionSettings = environment.configuration.languageVersionSettings
+            for (firFile in fir) {
+                val packageFragment = serializeSingleFirFile(
+                    firFile,
+                    session,
+                    scopeSession,
+                    FirKLibSerializerExtension(session, metadataVersion, FirElementAwareSerializableStringTable()),
+                    languageVersionSettings
+                )
+                fragments.getOrPut(firFile.packageFqName.asString()) { mutableListOf() }.add(packageFragment.toByteArray())
+            }
        }

        val header = KlibMetadataProtoBuf.Header.newBuilder()
-        header.moduleName = session.moduleData.name.asString()
+        header.moduleName = analysisResult.last().session.moduleData.name.asString()

        if (configuration.languageVersionSettings.isPreRelease()) {
            header.flags = KlibMetadataHeaderFlags.PRE_RELEASE