K1/K2: deduplicate/simplify the code around JVM CLI

Related to KT-61510, KT-61745
2023-09-06 14:36:05 +02:00
parent 3561054998
commit b0b6a6a05a
6 changed files with 281 additions and 322 deletions
@@ -139,7 +139,7 @@ class K2JVMCompiler : CLICompiler<K2JVMCompilerArguments>() {
            val targetDescription = chunk.map { input -> input.getModuleName() + "-" + input.getModuleType() }.let { names ->
                names.singleOrNull() ?: names.joinToString()
            }
-            // K2 works with multi-module chunks only in PSI mode
+            // K2 works with multi-module chunks only in PSI mode (KT-61745)
            if (chunk.size == 1 &&
                configuration.getBoolean(CommonConfigurationKeys.USE_FIR) &&
                configuration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE)
@@ -149,8 +149,13 @@ class K2JVMCompiler : CLICompiler<K2JVMCompilerArguments>() {
                    createProjectEnvironment(configuration, rootDisposable, EnvironmentConfigFiles.JVM_CONFIG_FILES, messageCollector)
                if (messageCollector.hasErrors()) return COMPILATION_ERROR

+                // TODO: uncomment (KT-61761)
+//                if (!FirKotlinToJvmBytecodeCompiler.checkNotSupportedPlugins(configuration, messageCollector)) {
+//                    return COMPILATION_ERROR
+//                }
+
                compileModulesUsingFrontendIrAndLightTree(
-                    projectEnvironment, configuration, messageCollector, buildFile, chunk, targetDescription,
+                    projectEnvironment, configuration, messageCollector, buildFile, chunk.single(), targetDescription,
                    checkSourceFiles = !arguments.allowNoSourceFiles && !arguments.version
                )
            } else {
@@ -9,16 +9,15 @@ package org.jetbrains.kotlin.cli.jvm.compiler

 import com.intellij.openapi.project.Project
 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.*
 import org.jetbrains.kotlin.cli.common.fir.FirDiagnosticsCompilerResultsReporter
 import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
-import org.jetbrains.kotlin.codegen.ClassBuilderFactories
-import org.jetbrains.kotlin.codegen.CodegenFactory
+import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.ModuleCompilerEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.ModuleCompilerIrBackendInput
+import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.generateCodeFromIr
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
@@ -27,6 +26,7 @@ import org.jetbrains.kotlin.constant.EvaluatedConstTracker
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
 import org.jetbrains.kotlin.diagnostics.impl.BaseDiagnosticsCollector
+import org.jetbrains.kotlin.diagnostics.impl.PendingDiagnosticsCollectorWithSuppress
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.backend.*
 import org.jetbrains.kotlin.fir.backend.jvm.*
@@ -35,8 +35,6 @@ import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.extensions.FirAnalysisHandlerExtension
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 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.types.arrayElementType
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.isArrayType
@@ -50,6 +48,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import org.jetbrains.kotlin.resolve.multiplatform.hmppModuleName
 import org.jetbrains.kotlin.resolve.multiplatform.isCommonSource
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
@@ -58,13 +57,13 @@ import java.io.File

 object FirKotlinToJvmBytecodeCompiler {
    fun checkNotSupportedPlugins(
-        projectConfiguration: CompilerConfiguration,
+        compilerConfiguration: CompilerConfiguration,
        messageCollector: MessageCollector
    ): Boolean {
        val notSupportedPlugins = mutableListOf<String?>().apply {
-            projectConfiguration.get(ComponentRegistrar.PLUGIN_COMPONENT_REGISTRARS)
+            compilerConfiguration.get(ComponentRegistrar.PLUGIN_COMPONENT_REGISTRARS)
                .collectIncompatiblePluginNamesTo(this, ComponentRegistrar::supportsK2)
-            projectConfiguration.get(CompilerPluginRegistrar.COMPILER_PLUGIN_REGISTRARS)
+            compilerConfiguration.get(CompilerPluginRegistrar.COMPILER_PLUGIN_REGISTRARS)
                .collectIncompatiblePluginNamesTo(this, CompilerPluginRegistrar::supportsK2)
        }

@@ -83,57 +82,47 @@ object FirKotlinToJvmBytecodeCompiler {
        return true
    }

-    fun compileModulesUsingFrontendIR(
-        projectEnvironment: AbstractProjectEnvironment,
-        projectConfiguration: CompilerConfiguration,
+    fun compileModulesUsingFrontendIRAndPsi(
+        projectEnvironment: VfsBasedProjectEnvironment,
+        compilerConfiguration: CompilerConfiguration,
        messageCollector: MessageCollector,
        allSources: List<KtFile>,
        buildFile: File?,
-        chunk: List<Module>
+        module: Module,
    ): Boolean {
-        val performanceManager = projectConfiguration.get(CLIConfigurationKeys.PERF_MANAGER)
+        val performanceManager = compilerConfiguration.get(CLIConfigurationKeys.PERF_MANAGER)

-        val outputs = ArrayList<Pair<FirResult, GenerationState>>(chunk.size)
-        val targetIds = projectConfiguration.get(JVMConfigurationKeys.MODULES)?.map(::TargetId)
-        val incrementalComponents = projectConfiguration.get(JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS)
-        val isMultiModuleChunk = chunk.size > 1 // Note: should be always false
+        val targetIds = compilerConfiguration.get(JVMConfigurationKeys.MODULES)?.map(::TargetId)
+        val incrementalComponents = compilerConfiguration.get(JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS)

-        // TODO: run lowerings for all modules in the chunk, then run codegen for all modules.
-        val project = (projectEnvironment as? VfsBasedProjectEnvironment)?.project
-        if (project != null) {
-            FirAnalysisHandlerExtension.analyze(project, projectConfiguration)?.let { return it }
-        }
+        val project = projectEnvironment.project
+        FirAnalysisHandlerExtension.analyze(project, compilerConfiguration)?.let { return it }

-        for (module in chunk) {
-            val moduleConfiguration = projectConfiguration.applyModuleProperties(module, buildFile)
-            val context = CompilationContext(
-                module,
-                module.getSourceFiles(
-                    allSources, (projectEnvironment as? VfsBasedProjectEnvironment)?.localFileSystem, isMultiModuleChunk, buildFile
-                ),
-                projectEnvironment,
-                messageCollector,
-                moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME),
-                moduleConfiguration,
-                performanceManager,
-                targetIds,
-                incrementalComponents,
-                extensionRegistrars = project?.let { FirExtensionRegistrar.getInstances(it) } ?: emptyList(),
-                irGenerationExtensions = project?.let { IrGenerationExtension.getInstances(it) } ?: emptyList()
-            )
-            val generationState = context.compileModule() ?: return false
-            outputs += generationState
-        }
+        val moduleConfiguration = compilerConfiguration.applyModuleProperties(module, buildFile)
+        val context = CompilationContext(
+            module,
+            allSources,
+            projectEnvironment,
+            messageCollector,
+            moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME),
+            moduleConfiguration,
+            performanceManager,
+            targetIds,
+            incrementalComponents,
+            extensionRegistrars = FirExtensionRegistrar.getInstances(project),
+            irGenerationExtensions = IrGenerationExtension.getInstances(project)
+        )
+        val resultAndGenerationState = context.compileModule() ?: return false

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

        return writeOutputsIfNeeded(
            project,
-            projectConfiguration,
-            chunk,
-            outputs.map(Pair<FirResult, GenerationState>::second),
+            compilerConfiguration,
+            messageCollector,
+            listOf(resultAndGenerationState.second),
            mainClassFqName
        )
    }
@@ -153,8 +142,8 @@ object FirKotlinToJvmBytecodeCompiler {
        if (!checkKotlinPackageUsageForPsi(configuration, allSources)) return null

        val renderDiagnosticNames = configuration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
+        val diagnosticsReporter = createPendingReporter(messageCollector)

-        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
        val firResult = runFrontend(allSources, diagnosticsReporter, module.getModuleName(), module.getFriendPaths()).also {
            performanceManager?.notifyAnalysisFinished()
        }
@@ -174,9 +163,8 @@ object FirKotlinToJvmBytecodeCompiler {
        performanceManager?.notifyIRTranslationFinished()

        val generationState = runBackend(
-            allSources,
-            fir2IrAndIrActualizerResult,
            fir2IrExtensions,
+            fir2IrAndIrActualizerResult,
            diagnosticsReporter
        )

@@ -189,6 +177,11 @@ object FirKotlinToJvmBytecodeCompiler {
        return firResult to generationState
    }

+    fun createPendingReporter(messageCollector: MessageCollector): PendingDiagnosticsCollectorWithSuppress =
+        DiagnosticReporterFactory.createPendingReporter { isError, message ->
+            messageCollector.report(if (isError) CompilerMessageSeverity.ERROR else CompilerMessageSeverity.WARNING, message)
+        }
+
    fun FrontendContext.runFrontend(
        ktFiles: List<KtFile>,
        diagnosticsReporter: BaseDiagnosticsCollector,
@@ -199,8 +192,7 @@ object FirKotlinToJvmBytecodeCompiler {
            AnalyzerWithCompilerReport.reportSyntaxErrors(ktFile, messageCollector).isHasErrors or errorsFound
        }

-        val sourceScope = (projectEnvironment as VfsBasedProjectEnvironment).getSearchScopeByPsiFiles(ktFiles) +
-                projectEnvironment.getSearchScopeForProjectJavaSources()
+        val sourceScope = projectEnvironment.getSearchScopeByPsiFiles(ktFiles) + projectEnvironment.getSearchScopeForProjectJavaSources()

        var librariesScope = projectEnvironment.getSearchScopeForProjectLibraries()

@@ -254,67 +246,46 @@ object FirKotlinToJvmBytecodeCompiler {
    }

    private fun CompilationContext.runBackend(
-        ktFiles: List<KtFile>,
+        fir2IrExtensions: JvmFir2IrExtensions,
        fir2IrActualizedResult: Fir2IrActualizedResult,
-        extensions: JvmGeneratorExtensions,
-        diagnosticsReporter: BaseDiagnosticsCollector
+        diagnosticsReporter: BaseDiagnosticsCollector,
    ): GenerationState {
        val (moduleFragment, components, pluginContext, irActualizedResult) = fir2IrActualizedResult
-        val dummyBindingContext = NoScopeRecordCliBindingTrace().bindingContext
-        val codegenFactory = JvmIrCodegenFactory(
+        val irInput = ModuleCompilerIrBackendInput(
+            TargetId(module),
            configuration,
-            configuration.get(CLIConfigurationKeys.PHASE_CONFIG),
+            fir2IrExtensions,
+            moduleFragment,
+            components,
+            pluginContext,
+            irActualizedResult
        )

-        val generationState = GenerationState.Builder(
-            (projectEnvironment as VfsBasedProjectEnvironment).project, ClassBuilderFactories.BINARIES,
-            moduleFragment.descriptor, dummyBindingContext, configuration
-        ).withModule(
-            module
-        ).onIndependentPartCompilationEnd(
-            createOutputFilesFlushingCallbackIfPossible(configuration)
-        ).isIrBackend(
-            true
-        ).jvmBackendClassResolver(
-            FirJvmBackendClassResolver(components)
-        ).diagnosticReporter(
-            diagnosticsReporter
-        ).build()
-
        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()

-        performanceManager?.notifyIRLoweringStarted()
-        generationState.beforeCompile()
-        generationState.oldBEInitTrace(ktFiles)
-        codegenFactory.generateModuleInFrontendIRMode(
-            generationState, moduleFragment, components.symbolTable, components.irProviders,
-            extensions, FirJvmBackendExtension(components, irActualizedResult), pluginContext
-        ) {
-            performanceManager?.notifyIRLoweringFinished()
-            performanceManager?.notifyIRGenerationStarted()
-        }
-        CodegenFactory.doCheckCancelled(generationState)
-        generationState.factory.done()
+        val generationState = generateCodeFromIr(
+            irInput, ModuleCompilerEnvironment(projectEnvironment, diagnosticsReporter), performanceManager
+        ).generationState

        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()

        AnalyzerWithCompilerReport.reportDiagnostics(
            FilteredJvmDiagnostics(
                generationState.collectedExtraJvmDiagnostics,
-                dummyBindingContext.diagnostics
+                Diagnostics.EMPTY
            ),
            messageCollector,
            renderDiagnosticName
        )
-        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()

+        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
        return generationState
    }

    private class CompilationContext(
        val module: Module,
        val allSources: List<KtFile>,
-        override val projectEnvironment: AbstractProjectEnvironment,
+        override val projectEnvironment: VfsBasedProjectEnvironment,
        override val messageCollector: MessageCollector,
        val renderDiagnosticName: Boolean,
        override val configuration: CompilerConfiguration,
@@ -326,7 +297,7 @@ object FirKotlinToJvmBytecodeCompiler {
    ) : FrontendContext

    class FrontendContextForMultiChunkMode private constructor(
-        override val projectEnvironment: AbstractProjectEnvironment,
+        override val projectEnvironment: VfsBasedProjectEnvironment,
        override val messageCollector: MessageCollector,
        override val incrementalComponents: IncrementalCompilationComponents?,
        override val extensionRegistrars: List<FirExtensionRegistrar>,
@@ -334,22 +305,22 @@ object FirKotlinToJvmBytecodeCompiler {
        override val targetIds: List<TargetId>?
    ) : FrontendContext {
        constructor(
-            projectEnvironment: AbstractProjectEnvironment,
+            projectEnvironment: VfsBasedProjectEnvironment,
            environment: KotlinCoreEnvironment,
-            projectConfiguration: CompilerConfiguration,
+            compilerConfiguration: CompilerConfiguration,
            project: Project?
        ) : this(
            projectEnvironment,
            environment.messageCollector,
-            incrementalComponents = projectConfiguration.get(JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS),
+            incrementalComponents = compilerConfiguration.get(JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS),
            extensionRegistrars = project?.let { FirExtensionRegistrar.getInstances(it) } ?: emptyList(),
-            configuration = projectConfiguration,
-            targetIds = projectConfiguration.get(JVMConfigurationKeys.MODULES)?.map(::TargetId)
+            configuration = compilerConfiguration,
+            targetIds = compilerConfiguration.get(JVMConfigurationKeys.MODULES)?.map(::TargetId)
        )
    }

    interface FrontendContext {
-        val projectEnvironment: AbstractProjectEnvironment
+        val projectEnvironment: VfsBasedProjectEnvironment
        val messageCollector: MessageCollector
        val incrementalComponents: IncrementalCompilationComponents?
        val extensionRegistrars: List<FirExtensionRegistrar>
@@ -50,6 +50,7 @@ import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStat
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade
+import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import java.io.File

 object KotlinToJVMBytecodeCompiler {
@@ -61,9 +62,10 @@ object KotlinToJVMBytecodeCompiler {
    ): Boolean {
        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()

-        val repeats = environment.configuration[CLIConfigurationKeys.REPEAT_COMPILE_MODULES]
+        val compilerConfiguration = environment.configuration
+        val repeats = compilerConfiguration[CLIConfigurationKeys.REPEAT_COMPILE_MODULES]
        if (repeats != null && !repeat) {
-            val performanceManager = environment.configuration[CLIConfigurationKeys.PERF_MANAGER]
+            val performanceManager = compilerConfiguration[CLIConfigurationKeys.PERF_MANAGER]
            return (0 until repeats).map {
                val result = compileModules(environment, buildFile, chunk, repeat = true)
                performanceManager?.notifyRepeat(repeats, it)
@@ -71,102 +73,49 @@ object KotlinToJVMBytecodeCompiler {
            }.last()
        }

-        val moduleVisibilityManager = ModuleVisibilityManager.SERVICE.getInstance(environment.project)
+        val project = environment.project
+        val moduleVisibilityManager = ModuleVisibilityManager.SERVICE.getInstance(project)
        for (module in chunk) {
            moduleVisibilityManager.addModule(module)
        }

-        val friendPaths = environment.configuration.getList(JVMConfigurationKeys.FRIEND_PATHS)
+        val friendPaths = compilerConfiguration.getList(JVMConfigurationKeys.FRIEND_PATHS)
        for (path in friendPaths) {
            moduleVisibilityManager.addFriendPath(path)
        }

-        val projectConfiguration = environment.configuration
-        var mainClassFqName: FqName? = null
-        val moduleDescriptor: ModuleDescriptor
-        val bindingContext: BindingContext
-        var firJvmBackendClassResolver: FirJvmBackendClassResolver? = null
-        var firJvmBackendExtension: FirJvmBackendExtension? = null
-        val irResult = if (projectConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR)) {
+        val useFrontendIR = compilerConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR)
+        val messageCollector = environment.messageCollector
+        val (codegenFactory, wholeBackendInput, moduleDescriptor, bindingContext, firJvmBackendResolver, firJvmBackendExtension, mainClassFqName) = if (useFrontendIR) {
            // K2/PSI: base checks
            val projectEnvironment =
                VfsBasedProjectEnvironment(
-                    environment.project,
+                    project,
                    VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
                ) { environment.createPackagePartProvider(it) }

-            if (!FirKotlinToJvmBytecodeCompiler.checkNotSupportedPlugins(
-                    projectConfiguration, environment.messageCollector
-                )
-            ) return false
+            if (!FirKotlinToJvmBytecodeCompiler.checkNotSupportedPlugins(compilerConfiguration, messageCollector)) {
+                return false
+            }

-            // K2/PSI: single chunk mode fallback
+            // K2/PSI: single module chunk mode fallback (KT-61745)
            if (chunk.size == 1) {
-                return FirKotlinToJvmBytecodeCompiler.compileModulesUsingFrontendIR(
+                return FirKotlinToJvmBytecodeCompiler.compileModulesUsingFrontendIRAndPsi(
                    projectEnvironment,
-                    environment.configuration,
-                    environment.messageCollector,
+                    compilerConfiguration,
+                    messageCollector,
                    environment.getSourceFiles(),
-                    buildFile, chunk
+                    buildFile,
+                    chunk.single()
                )
            }

-            val sourceFiles = environment.getSourceFiles()
-            val project = (projectEnvironment as? VfsBasedProjectEnvironment)?.project
-            val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
-            val frontendContext = FirKotlinToJvmBytecodeCompiler.FrontendContextForMultiChunkMode(
-                projectEnvironment, environment, projectConfiguration, project
-            )
-
-            with(frontendContext) {
-                // K2/PSI: frontend
-                val firResult = runFrontend(
-                    sourceFiles, diagnosticsReporter, chunk.joinToString(separator = "+") { it.getModuleName() },
-                    chunk.fold(emptyList()) { paths, m -> paths + m.getFriendPaths() }
-                ) ?: run {
-                    FirDiagnosticsCompilerResultsReporter.reportToMessageCollector(
-                        diagnosticsReporter, messageCollector,
-                        projectConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
-                    )
-                    return false
-                }
-                // K2/PSI: FIR2IR
-                val fir2IrExtensions = JvmFir2IrExtensions(configuration, JvmIrDeserializerImpl(), JvmIrMangler)
-                val irGenerationExtensions = project?.let { IrGenerationExtension.getInstances(it) } ?: emptyList()
-                val fir2IrAndIrActualizerResult = convertToIrAndActualizeForJvm(
-                    firResult, diagnosticsReporter, fir2IrExtensions, irGenerationExtensions
-                )
-                moduleDescriptor = fir2IrAndIrActualizerResult.irModuleFragment.descriptor
-                bindingContext = NoScopeRecordCliBindingTrace().bindingContext
-                firJvmBackendClassResolver = FirJvmBackendClassResolver(fir2IrAndIrActualizerResult.components)
-                firJvmBackendExtension = FirJvmBackendExtension(
-                    fir2IrAndIrActualizerResult.components, fir2IrAndIrActualizerResult.irActualizedResult
-                )
-                fir2IrAndIrActualizerResult.codegenFactoryWithJvmIrBackendInput(configuration)
-            }
+            runFrontendAndGenerateIrForMultiModuleChunkUsingFrontendIR(environment, projectEnvironment, compilerConfiguration, chunk)
        } else {
-            // K1: Frontend
-            val result = repeatAnalysisIfNeeded(analyze(environment), environment)
-            if (result == null || !result.shouldGenerateCode) return false
-            moduleDescriptor = result.moduleDescriptor
-            bindingContext = result.bindingContext
-
-            ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
-
-            result.throwIfError()
-
-            if (chunk.size == 1 && projectConfiguration.get(JVMConfigurationKeys.OUTPUT_JAR) != null) {
-                mainClassFqName = findMainClass(
-                    result.bindingContext, projectConfiguration.languageVersionSettings, environment.getSourceFiles()
-                )
-            }
-
-            // K1: PSI2IR
-            convertToIr(environment, result)
-        }
+            runFrontendAndGenerateIrUsingClassicFrontend(environment, compilerConfiguration, chunk)
+        } ?: return false
        // K1/K2 common multi-chunk part
        val localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
-        val (codegenFactory, wholeBackendInput) = irResult

        val diagnosticsReporter = DiagnosticReporterFactory.createReporter()

@@ -176,18 +125,18 @@ object KotlinToJVMBytecodeCompiler {
            ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()

            val ktFiles = module.getSourceFiles(environment.getSourceFiles(), localFileSystem, chunk.size > 1, buildFile)
-            if (!checkKotlinPackageUsageForPsi(environment.configuration, ktFiles)) return false
-            val moduleConfiguration = projectConfiguration.applyModuleProperties(module, buildFile)
+            if (!checkKotlinPackageUsageForPsi(compilerConfiguration, ktFiles)) return false
+            val moduleConfiguration = compilerConfiguration.applyModuleProperties(module, buildFile)

            val backendInput = codegenFactory.getModuleChunkBackendInput(wholeBackendInput, ktFiles).let {
                if (it is JvmIrCodegenFactory.JvmIrBackendInput && firJvmBackendExtension != null) {
-                    it.copy(backendExtension = firJvmBackendExtension!!)
+                    it.copy(backendExtension = firJvmBackendExtension)
                } else it
            }
            // Lowerings (per module)
            codegenInputs += runLowerings(
                environment, moduleConfiguration, moduleDescriptor, bindingContext,
-                ktFiles, module, codegenFactory, backendInput, diagnosticsReporter, firJvmBackendClassResolver
+                ktFiles, module, codegenFactory, backendInput, diagnosticsReporter, firJvmBackendResolver
            )
        }

@@ -195,12 +144,97 @@ object KotlinToJVMBytecodeCompiler {

        for (input in codegenInputs) {
            // Codegen (per module)
-            outputs += runCodegen(input, input.state, codegenFactory, bindingContext, diagnosticsReporter, environment.configuration)
+            outputs += runCodegen(input, input.state, codegenFactory, bindingContext, diagnosticsReporter, compilerConfiguration)
        }

-        return writeOutputsIfNeeded(environment.project, projectConfiguration, chunk, outputs, mainClassFqName)
+        return writeOutputsIfNeeded(project, compilerConfiguration, messageCollector, outputs, mainClassFqName)
    }

+    private fun runFrontendAndGenerateIrForMultiModuleChunkUsingFrontendIR(
+        environment: KotlinCoreEnvironment,
+        projectEnvironment: VfsBasedProjectEnvironment,
+        compilerConfiguration: CompilerConfiguration,
+        chunk: List<Module>,
+    ): BackendInputForMultiModuleChunk? {
+        val sourceFiles = environment.getSourceFiles()
+        val project = projectEnvironment.project
+        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
+        val frontendContext = FirKotlinToJvmBytecodeCompiler.FrontendContextForMultiChunkMode(
+            projectEnvironment, environment, compilerConfiguration, project
+        )
+
+        with(frontendContext) {
+            // K2/PSI: frontend
+            val firResult = runFrontend(
+                sourceFiles, diagnosticsReporter, chunk.joinToString(separator = "+") { it.getModuleName() },
+                chunk.fold(emptyList()) { paths, m -> paths + m.getFriendPaths() }
+            ) ?: run {
+                FirDiagnosticsCompilerResultsReporter.reportToMessageCollector(
+                    diagnosticsReporter, messageCollector,
+                    compilerConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
+                )
+                return null
+            }
+            // K2/PSI: FIR2IR
+            val fir2IrExtensions = JvmFir2IrExtensions(configuration, JvmIrDeserializerImpl(), JvmIrMangler)
+            val irGenerationExtensions = IrGenerationExtension.getInstances(project)
+            val fir2IrAndIrActualizerResult = convertToIrAndActualizeForJvm(
+                firResult, diagnosticsReporter, fir2IrExtensions, irGenerationExtensions
+            )
+            val (factory, input) = fir2IrAndIrActualizerResult.codegenFactoryWithJvmIrBackendInput(configuration)
+            return BackendInputForMultiModuleChunk(
+                factory,
+                input,
+                fir2IrAndIrActualizerResult.irModuleFragment.descriptor,
+                NoScopeRecordCliBindingTrace().bindingContext,
+                FirJvmBackendClassResolver(fir2IrAndIrActualizerResult.components),
+                FirJvmBackendExtension(
+                    fir2IrAndIrActualizerResult.components, fir2IrAndIrActualizerResult.irActualizedResult
+                )
+            )
+        }
+    }
+
+    private fun runFrontendAndGenerateIrUsingClassicFrontend(
+        environment: KotlinCoreEnvironment,
+        compilerConfiguration: CompilerConfiguration,
+        chunk: List<Module>
+    ): BackendInputForMultiModuleChunk? {
+        // K1: Frontend
+        val result = repeatAnalysisIfNeeded(analyze(environment), environment)
+        if (result == null || !result.shouldGenerateCode) return null
+
+        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+
+        result.throwIfError()
+
+        val mainClassFqName = runIf(chunk.size == 1 && compilerConfiguration.get(JVMConfigurationKeys.OUTPUT_JAR) != null) {
+            findMainClass(
+                result.bindingContext, compilerConfiguration.languageVersionSettings, environment.getSourceFiles()
+            )
+        }
+
+        // K1: PSI2IR
+        val (factory, input) = convertToIr(environment, result)
+        return BackendInputForMultiModuleChunk(
+            factory,
+            input,
+            result.moduleDescriptor,
+            result.bindingContext,
+            mainClassFqName = mainClassFqName
+        )
+    }
+
+    private data class BackendInputForMultiModuleChunk(
+        val codegenFactory: CodegenFactory,
+        val backendInput: CodegenFactory.BackendInput,
+        val moduleDescriptor: ModuleDescriptor,
+        val bindingContext: BindingContext,
+        val firJvmBackendClassResolver: FirJvmBackendClassResolver? = null,
+        val firJvmBackendExtension: FirJvmBackendExtension? = null,
+        val mainClassFqName: FqName? = null,
+    )
+
    fun compileBunchOfSources(environment: KotlinCoreEnvironment): Boolean {
        val moduleVisibilityManager = ModuleVisibilityManager.SERVICE.getInstance(environment.project)

@@ -135,13 +135,13 @@ fun writeOutput(
 }

 fun writeOutputsIfNeeded(
-    project: Project?,
-    projectConfiguration: CompilerConfiguration,
-    chunk: List<Module>,
-    outputs: List<GenerationState>,
+    project: Project,
+    configuration: CompilerConfiguration,
+    messageCollector: MessageCollector,
+    outputs: Collection<GenerationState>,
    mainClassFqName: FqName?
 ): Boolean {
-    if (projectConfiguration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY).hasErrors()) {
+    if (messageCollector.hasErrors()) {
        return false
    }

@@ -154,16 +154,16 @@ fun writeOutputsIfNeeded(
        outputs.forEach(GenerationState::destroy)
    }

-    if (projectConfiguration.getBoolean(JVMConfigurationKeys.COMPILE_JAVA)) {
-        val singleModule = chunk.singleOrNull()
-        if (singleModule != null) {
-            return JavacWrapper.getInstance(project!!).use {
-                it.compile(File(singleModule.getOutputDirectory()))
+    if (configuration.getBoolean(JVMConfigurationKeys.COMPILE_JAVA)) {
+        val singleState = outputs.singleOrNull()
+        if (singleState != null) {
+            return JavacWrapper.getInstance(project).use {
+                it.compile(singleState.outDirectory)
            }
        } else {
-            projectConfiguration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY).report(
+            messageCollector.report(
                CompilerMessageSeverity.WARNING,
-                "A chunk contains multiple modules (${chunk.joinToString { it.getModuleName() }}). " +
+                "A chunk contains multiple modules (${outputs.joinToString { it.moduleName }}). " +
                        "-Xuse-javac option couldn't be used to compile java files"
            )
        }
@@ -41,7 +41,6 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.constant.EvaluatedConstTracker
-import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
 import org.jetbrains.kotlin.diagnostics.impl.BaseDiagnosticsCollector
 import org.jetbrains.kotlin.fir.backend.Fir2IrConfiguration
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmBackendClassResolver
@@ -57,124 +56,111 @@ 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.ir.backend.jvm.serialization.JvmIrMangler
-import org.jetbrains.kotlin.javac.JavacWrapper
 import org.jetbrains.kotlin.load.kotlin.MetadataFinderFactory
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.load.kotlin.VirtualFileFinderFactory
 import org.jetbrains.kotlin.load.kotlin.incremental.IncrementalPackagePartProvider
 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.utils.addToStdlib.runIf
 import java.io.File

 fun compileModulesUsingFrontendIrAndLightTree(
-    projectEnvironment: AbstractProjectEnvironment,
+    projectEnvironment: VfsBasedProjectEnvironment,
    compilerConfiguration: CompilerConfiguration,
    messageCollector: MessageCollector,
    buildFile: File?,
-    chunk: List<Module>,
+    module: Module,
    targetDescription: String,
    checkSourceFiles: Boolean
 ): Boolean {
-    require(projectEnvironment is VfsBasedProjectEnvironment) // TODO: abstract away this requirement
    ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()

    val performanceManager = compilerConfiguration[CLIConfigurationKeys.PERF_MANAGER]
-
    performanceManager?.notifyCompilerInitialized(0, 0, targetDescription)

    val project = projectEnvironment.project
    FirAnalysisHandlerExtension.analyze(project, compilerConfiguration)?.let { return it }

-    val outputs = mutableListOf<GenerationState>()
-    var mainClassFqName: FqName? = null
-
-    for (module in chunk) {
-        val moduleConfiguration = compilerConfiguration.copy().applyModuleProperties(module, buildFile).apply {
-            put(JVMConfigurationKeys.FRIEND_PATHS, module.getFriendPaths())
-        }
-        val groupedSources = collectSources(compilerConfiguration, projectEnvironment, messageCollector)
-        if (messageCollector.hasErrors()) {
-            return false
-        }
-
-        if (checkSourceFiles && groupedSources.isEmpty() && buildFile == null) {
-            messageCollector.report(CompilerMessageSeverity.ERROR, "No source files")
-            return false
-        }
-
-        val compilerInput = ModuleCompilerInput(
-            TargetId(module),
-            groupedSources,
-            CommonPlatforms.defaultCommonPlatform,
-            JvmPlatforms.unspecifiedJvmPlatform,
-            moduleConfiguration
-        )
-
-        val renderDiagnosticName = moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
-        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter { isError, message ->
-            messageCollector.report(if (isError) CompilerMessageSeverity.ERROR else CompilerMessageSeverity.WARNING, message)
-        }
-        val compilerEnvironment = ModuleCompilerEnvironment(projectEnvironment, diagnosticsReporter)
-
-        performanceManager?.notifyAnalysisStarted()
-
-        val analysisResults = compileModuleToAnalyzedFir(
-            compilerInput,
-            compilerEnvironment,
-            emptyList(),
-            null,
-            diagnosticsReporter,
-            performanceManager
-        )
-
-        if (!checkKotlinPackageUsageForLightTree(moduleConfiguration, analysisResults.outputs.flatMap { it.fir })) {
-            return false
-        }
-
-        performanceManager?.notifyAnalysisFinished()
-
-        // TODO: consider what to do if many modules has main classes
-        if (mainClassFqName == null && moduleConfiguration.get(JVMConfigurationKeys.OUTPUT_JAR) != null) {
-            mainClassFqName = findMainClass(analysisResults.outputs.last().fir)
-        }
-
-        if (diagnosticsReporter.hasErrors) {
-            diagnosticsReporter.reportToMessageCollector(messageCollector, renderDiagnosticName)
-            continue
-        }
-
-        performanceManager?.notifyGenerationStarted()
-        performanceManager?.notifyIRTranslationStarted()
-
-        val irInput = convertAnalyzedFirToIr(compilerInput, analysisResults, compilerEnvironment)
-
-        performanceManager?.notifyIRTranslationFinished()
-
-        val codegenOutput = generateCodeFromIr(irInput, compilerEnvironment, performanceManager)
-
-        diagnosticsReporter.reportToMessageCollector(
-            messageCollector, moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
-        )
-
-        performanceManager?.notifyIRGenerationFinished()
-        performanceManager?.notifyGenerationFinished()
-
-        if (!diagnosticsReporter.hasErrors) {
-            outputs.add(codegenOutput.generationState)
-        }
+    val moduleConfiguration = compilerConfiguration.copy().applyModuleProperties(module, buildFile).apply {
+        put(JVMConfigurationKeys.FRIEND_PATHS, module.getFriendPaths())
+    }
+    val groupedSources = collectSources(compilerConfiguration, projectEnvironment, messageCollector)
+    if (messageCollector.hasErrors()) {
+        return false
    }

-    return writeOutputs(
+    if (checkSourceFiles && groupedSources.isEmpty() && buildFile == null) {
+        messageCollector.report(CompilerMessageSeverity.ERROR, "No source files")
+        return false
+    }
+
+    val compilerInput = ModuleCompilerInput(
+        TargetId(module),
+        groupedSources,
+        CommonPlatforms.defaultCommonPlatform,
+        JvmPlatforms.unspecifiedJvmPlatform,
+        moduleConfiguration
+    )
+
+    val renderDiagnosticNames = moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
+    val diagnosticsReporter = FirKotlinToJvmBytecodeCompiler.createPendingReporter(messageCollector)
+
+
+    performanceManager?.notifyAnalysisStarted()
+
+    val analysisResults = compileModuleToAnalyzedFir(
+        compilerInput,
        projectEnvironment,
+        emptyList(),
+        null,
+        diagnosticsReporter,
+        performanceManager
+    )
+
+    if (!checkKotlinPackageUsageForLightTree(moduleConfiguration, analysisResults.outputs.flatMap { it.fir })) {
+        return false
+    }
+
+    performanceManager?.notifyAnalysisFinished()
+
+    val mainClassFqName = runIf(moduleConfiguration.get(JVMConfigurationKeys.OUTPUT_JAR) != null) {
+        findMainClass(analysisResults.outputs.last().fir)
+    }
+
+    if (diagnosticsReporter.hasErrors) {
+        diagnosticsReporter.reportToMessageCollector(messageCollector, renderDiagnosticNames)
+        return false
+    }
+
+    performanceManager?.notifyGenerationStarted()
+    performanceManager?.notifyIRTranslationStarted()
+
+    val compilerEnvironment = ModuleCompilerEnvironment(projectEnvironment, diagnosticsReporter)
+    val irInput = convertAnalyzedFirToIr(compilerInput, analysisResults, compilerEnvironment)
+
+    performanceManager?.notifyIRTranslationFinished()
+
+    val codegenOutput = generateCodeFromIr(irInput, compilerEnvironment, performanceManager)
+
+    diagnosticsReporter.reportToMessageCollector(
+        messageCollector, moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
+    )
+
+    performanceManager?.notifyIRGenerationFinished()
+    performanceManager?.notifyGenerationFinished()
+
+    return writeOutputsIfNeeded(
+        project,
        compilerConfiguration,
-        outputs,
+        messageCollector,
+        listOf(codegenOutput.generationState),
        mainClassFqName
    )
 }
@@ -200,7 +186,7 @@ fun convertAnalyzedFirToIr(
        inlineConstTracker = input.configuration[CommonConfigurationKeys.INLINE_CONST_TRACKER],
        allowNonCachedDeclarations = false,
    )
-    val (irModuleFragment, components, pluginContext, irActualizedResult) =
+    val (moduleFragment, components, pluginContext, irActualizedResult) =
        analysisResults.convertToIrAndActualizeForJvm(
            extensions, fir2IrConfiguration, irGenerationExtensions,
        )
@@ -209,7 +195,7 @@ fun convertAnalyzedFirToIr(
        input.targetId,
        input.configuration,
        extensions,
-        irModuleFragment,
+        moduleFragment,
        components,
        pluginContext,
        irActualizedResult
@@ -248,7 +234,6 @@ fun generateCodeFromIr(
    ).build()

    performanceManager?.notifyIRLoweringStarted()
-
    generationState.beforeCompile()
    codegenFactory.generateModuleInFrontendIRMode(
        generationState,
@@ -270,13 +255,12 @@ fun generateCodeFromIr(

 fun compileModuleToAnalyzedFir(
    input: ModuleCompilerInput,
-    environment: ModuleCompilerEnvironment,
+    projectEnvironment: VfsBasedProjectEnvironment,
    previousStepsSymbolProviders: List<FirSymbolProvider>,
    incrementalExcludesScope: AbstractProjectFileSearchScope?,
    diagnosticsReporter: BaseDiagnosticsCollector,
    performanceManager: CommonCompilerPerformanceManager?
 ): FirResult {
-    val projectEnvironment = environment.projectEnvironment
    val moduleConfiguration = input.configuration

    var librariesScope = projectEnvironment.getSearchScopeForProjectLibraries()
@@ -324,39 +308,6 @@ fun compileModuleToAnalyzedFir(
    return FirResult(outputs)
 }

-fun writeOutputs(
-    projectEnvironment: AbstractProjectEnvironment,
-    configuration: CompilerConfiguration,
-    outputs: Collection<GenerationState>,
-    mainClassFqName: FqName?
-): Boolean {
-    try {
-        for (state in outputs) {
-            ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
-            writeOutput(state.configuration, state.factory, mainClassFqName)
-        }
-    } finally {
-        outputs.forEach(GenerationState::destroy)
-    }
-
-    if (configuration.getBoolean(JVMConfigurationKeys.COMPILE_JAVA)) {
-        val singleState = outputs.singleOrNull()
-        if (singleState != null) {
-            return JavacWrapper.getInstance((projectEnvironment as VfsBasedProjectEnvironment).project).use {
-                it.compile(singleState.outDirectory)
-            }
-        } else {
-            configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY).report(
-                CompilerMessageSeverity.WARNING,
-                "A chunk contains multiple modules (${outputs.joinToString { it.moduleName }}). " +
-                        "-Xuse-javac option couldn't be used to compile java files"
-            )
-        }
-    }
-
-    return true
-}
-
 fun createIncrementalCompilationScope(
    configuration: CompilerConfiguration,
    projectEnvironment: AbstractProjectEnvironment,
@@ -370,7 +321,7 @@ fun createIncrementalCompilationScope(
    }
 }

-private fun needCreateIncrementalCompilationScope(configuration: CompilerConfiguration, ): Boolean {
+private fun needCreateIncrementalCompilationScope(configuration: CompilerConfiguration): Boolean {
    if (configuration.get(JVMConfigurationKeys.MODULES) == null) return false
    if (configuration.get(JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS) == null) return false
    return true
@@ -29,10 +29,7 @@ import org.jetbrains.kotlin.cli.common.messages.IrMessageCollector
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.cli.common.modules.ModuleBuilder
 import org.jetbrains.kotlin.cli.jvm.*
-import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
-import org.jetbrains.kotlin.cli.jvm.compiler.VfsBasedProjectEnvironment
-import org.jetbrains.kotlin.cli.jvm.compiler.findMainClass
-import org.jetbrains.kotlin.cli.jvm.compiler.forAllFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.*
 import org.jetbrains.kotlin.cli.jvm.compiler.pipeline.*
 import org.jetbrains.kotlin.cli.jvm.config.*
 import org.jetbrains.kotlin.cli.jvm.plugins.PluginCliParser
@@ -221,7 +218,7 @@ open class IncrementalFirJvmCompilerRunner(
                    val analysisResults =
                        compileModuleToAnalyzedFir(
                            compilerInput,
-                            compilerEnvironment,
+                            projectEnvironment,
                            emptyList(),
                            incrementalExcludesScope,
                            diagnosticsReporter,
@@ -308,9 +305,10 @@ open class IncrementalFirJvmCompilerRunner(

            diagnosticsReporter.reportToMessageCollector(messageCollector, renderDiagnosticName)

-            writeOutputs(
-                projectEnvironment,
+            writeOutputsIfNeeded(
+                projectEnvironment.project,
                configuration,
+                messageCollector,
                listOf(codegenOutput.generationState),
                mainClassFqName
            )