K1/K2: partially share code for multi-module chunks building

#KT-61510 Fixed Related to KT-61716
2023-09-05 11:48:07 +02:00
parent a1c2ded8a3
commit 3561054998
3 changed files with 231 additions and 73 deletions
@@ -139,7 +139,11 @@ class K2JVMCompiler : CLICompiler<K2JVMCompilerArguments>() {
            val targetDescription = chunk.map { input -> input.getModuleName() + "-" + input.getModuleType() }.let { names ->
                names.singleOrNull() ?: names.joinToString()
            }
-            if (configuration.getBoolean(CommonConfigurationKeys.USE_FIR) && configuration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE))  {
+            // K2 works with multi-module chunks only in PSI mode
+            if (chunk.size == 1 &&
+                configuration.getBoolean(CommonConfigurationKeys.USE_FIR) &&
+                configuration.getBoolean(CommonConfigurationKeys.USE_LIGHT_TREE)
+            ) {
                if (messageCollector.hasErrors()) return COMPILATION_ERROR
                val projectEnvironment =
                    createProjectEnvironment(configuration, rootDisposable, EnvironmentConfigFiles.JVM_CONFIG_FILES, messageCollector)
@@ -7,6 +7,7 @@

 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
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
 import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
 import org.jetbrains.kotlin.config.*
 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.fir.*
@@ -55,16 +57,10 @@ import org.jetbrains.kotlin.utils.addToStdlib.runUnless
 import java.io.File

 object FirKotlinToJvmBytecodeCompiler {
-    fun compileModulesUsingFrontendIR(
-        projectEnvironment: AbstractProjectEnvironment,
+    fun checkNotSupportedPlugins(
        projectConfiguration: CompilerConfiguration,
-        messageCollector: MessageCollector,
-        allSources: List<KtFile>,
-        buildFile: File?,
-        chunk: List<Module>
+        messageCollector: MessageCollector
    ): Boolean {
-        val performanceManager = projectConfiguration.get(CLIConfigurationKeys.PERF_MANAGER)
-
        val notSupportedPlugins = mutableListOf<String?>().apply {
            projectConfiguration.get(ComponentRegistrar.PLUGIN_COMPONENT_REGISTRARS)
                .collectIncompatiblePluginNamesTo(this, ComponentRegistrar::supportsK2)
@@ -84,10 +80,23 @@ object FirKotlinToJvmBytecodeCompiler {
            return false
        }

+        return true
+    }
+
+    fun compileModulesUsingFrontendIR(
+        projectEnvironment: AbstractProjectEnvironment,
+        projectConfiguration: CompilerConfiguration,
+        messageCollector: MessageCollector,
+        allSources: List<KtFile>,
+        buildFile: File?,
+        chunk: List<Module>
+    ): Boolean {
+        val performanceManager = projectConfiguration.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
+        val isMultiModuleChunk = chunk.size > 1 // Note: should be always false

        // TODO: run lowerings for all modules in the chunk, then run codegen for all modules.
        val project = (projectEnvironment as? VfsBasedProjectEnvironment)?.project
@@ -141,12 +150,12 @@ object FirKotlinToJvmBytecodeCompiler {
        performanceManager?.notifyAnalysisStarted()
        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()

-        if (!checkKotlinPackageUsageForPsi(moduleConfiguration, allSources)) return null
+        if (!checkKotlinPackageUsageForPsi(configuration, allSources)) return null

-        val renderDiagnosticNames = moduleConfiguration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
+        val renderDiagnosticNames = configuration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)

        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
-        val firResult = runFrontend(allSources, diagnosticsReporter).also {
+        val firResult = runFrontend(allSources, diagnosticsReporter, module.getModuleName(), module.getFriendPaths()).also {
            performanceManager?.notifyAnalysisFinished()
        }
        if (firResult == null) {
@@ -157,20 +166,9 @@ object FirKotlinToJvmBytecodeCompiler {
        performanceManager?.notifyGenerationStarted()
        performanceManager?.notifyIRTranslationStarted()

-        val fir2IrExtensions = JvmFir2IrExtensions(moduleConfiguration, JvmIrDeserializerImpl(), JvmIrMangler)
-        val fir2IrConfiguration = Fir2IrConfiguration(
-            languageVersionSettings = moduleConfiguration.languageVersionSettings,
-            diagnosticReporter = diagnosticsReporter,
-            linkViaSignatures = moduleConfiguration.getBoolean(JVMConfigurationKeys.LINK_VIA_SIGNATURES),
-            evaluatedConstTracker = moduleConfiguration
-                .putIfAbsent(CommonConfigurationKeys.EVALUATED_CONST_TRACKER, EvaluatedConstTracker.create()),
-            inlineConstTracker = moduleConfiguration[CommonConfigurationKeys.INLINE_CONST_TRACKER],
-            allowNonCachedDeclarations = false,
-        )
-        val fir2IrAndIrActualizerResult = firResult.convertToIrAndActualizeForJvm(
-            fir2IrExtensions,
-            fir2IrConfiguration,
-            irGenerationExtensions,
+        val fir2IrExtensions = JvmFir2IrExtensions(configuration, JvmIrDeserializerImpl(), JvmIrMangler)
+        val fir2IrAndIrActualizerResult = convertToIrAndActualizeForJvm(
+            firResult, diagnosticsReporter, fir2IrExtensions, irGenerationExtensions
        )

        performanceManager?.notifyIRTranslationFinished()
@@ -191,7 +189,12 @@ object FirKotlinToJvmBytecodeCompiler {
        return firResult to generationState
    }

-    private fun CompilationContext.runFrontend(ktFiles: List<KtFile>, diagnosticsReporter: BaseDiagnosticsCollector): FirResult? {
+    fun FrontendContext.runFrontend(
+        ktFiles: List<KtFile>,
+        diagnosticsReporter: BaseDiagnosticsCollector,
+        rootModuleName: String,
+        friendPaths: List<String>,
+    ): FirResult? {
        val syntaxErrors = ktFiles.fold(false) { errorsFound, ktFile ->
            AnalyzerWithCompilerReport.reportSyntaxErrors(ktFile, messageCollector).isHasErrors or errorsFound
        }
@@ -204,7 +207,7 @@ object FirKotlinToJvmBytecodeCompiler {
        val providerAndScopeForIncrementalCompilation = createContextForIncrementalCompilation(
            projectEnvironment,
            incrementalComponents,
-            moduleConfiguration,
+            configuration,
            targetIds,
            sourceScope
        )
@@ -212,10 +215,9 @@ object FirKotlinToJvmBytecodeCompiler {
        providerAndScopeForIncrementalCompilation?.precompiledBinariesFileScope?.let {
            librariesScope -= it
        }
-        val rootModuleName = module.getModuleName()
-        val libraryList = createLibraryListForJvm(rootModuleName, moduleConfiguration, module.getFriendPaths())
+        val libraryList = createLibraryListForJvm(rootModuleName, configuration, friendPaths)
        val sessionsWithSources = prepareJvmSessions(
-            ktFiles, moduleConfiguration, projectEnvironment, Name.special("<$rootModuleName>"),
+            ktFiles, configuration, projectEnvironment, Name.special("<$rootModuleName>"),
            extensionRegistrars, librariesScope, libraryList,
            isCommonSource = { it.isCommonSource == true },
            fileBelongsToModule = { file, moduleName -> file.hmppModuleName == moduleName },
@@ -229,6 +231,28 @@ object FirKotlinToJvmBytecodeCompiler {
        return runUnless(syntaxErrors || diagnosticsReporter.hasErrors) { FirResult(outputs) }
    }

+    fun FrontendContext.convertToIrAndActualizeForJvm(
+        firResult: FirResult,
+        diagnosticsReporter: DiagnosticReporter,
+        fir2IrExtensions: JvmFir2IrExtensions,
+        irGenerationExtensions: Collection<IrGenerationExtension>,
+    ): Fir2IrActualizedResult {
+        val fir2IrConfiguration = Fir2IrConfiguration(
+            languageVersionSettings = configuration.languageVersionSettings,
+            diagnosticReporter = diagnosticsReporter,
+            linkViaSignatures = configuration.getBoolean(JVMConfigurationKeys.LINK_VIA_SIGNATURES),
+            evaluatedConstTracker = configuration
+                .putIfAbsent(CommonConfigurationKeys.EVALUATED_CONST_TRACKER, EvaluatedConstTracker.create()),
+            inlineConstTracker = configuration[CommonConfigurationKeys.INLINE_CONST_TRACKER],
+            allowNonCachedDeclarations = false,
+        )
+        return firResult.convertToIrAndActualizeForJvm(
+            fir2IrExtensions,
+            fir2IrConfiguration,
+            irGenerationExtensions,
+        )
+    }
+
    private fun CompilationContext.runBackend(
        ktFiles: List<KtFile>,
        fir2IrActualizedResult: Fir2IrActualizedResult,
@@ -238,17 +262,17 @@ object FirKotlinToJvmBytecodeCompiler {
        val (moduleFragment, components, pluginContext, irActualizedResult) = fir2IrActualizedResult
        val dummyBindingContext = NoScopeRecordCliBindingTrace().bindingContext
        val codegenFactory = JvmIrCodegenFactory(
-            moduleConfiguration,
-            moduleConfiguration.get(CLIConfigurationKeys.PHASE_CONFIG),
+            configuration,
+            configuration.get(CLIConfigurationKeys.PHASE_CONFIG),
        )

        val generationState = GenerationState.Builder(
            (projectEnvironment as VfsBasedProjectEnvironment).project, ClassBuilderFactories.BINARIES,
-            moduleFragment.descriptor, dummyBindingContext, moduleConfiguration
+            moduleFragment.descriptor, dummyBindingContext, configuration
        ).withModule(
            module
        ).onIndependentPartCompilationEnd(
-            createOutputFilesFlushingCallbackIfPossible(moduleConfiguration)
+            createOutputFilesFlushingCallbackIfPossible(configuration)
        ).isIrBackend(
            true
        ).jvmBackendClassResolver(
@@ -290,16 +314,48 @@ object FirKotlinToJvmBytecodeCompiler {
    private class CompilationContext(
        val module: Module,
        val allSources: List<KtFile>,
-        val projectEnvironment: AbstractProjectEnvironment,
-        val messageCollector: MessageCollector,
+        override val projectEnvironment: AbstractProjectEnvironment,
+        override val messageCollector: MessageCollector,
        val renderDiagnosticName: Boolean,
-        val moduleConfiguration: CompilerConfiguration,
+        override val configuration: CompilerConfiguration,
        val performanceManager: CommonCompilerPerformanceManager?,
-        val targetIds: List<TargetId>?,
-        val incrementalComponents: IncrementalCompilationComponents?,
-        val extensionRegistrars: List<FirExtensionRegistrar>,
+        override val targetIds: List<TargetId>?,
+        override val incrementalComponents: IncrementalCompilationComponents?,
+        override val extensionRegistrars: List<FirExtensionRegistrar>,
        val irGenerationExtensions: Collection<IrGenerationExtension>
-    )
+    ) : FrontendContext
+
+    class FrontendContextForMultiChunkMode private constructor(
+        override val projectEnvironment: AbstractProjectEnvironment,
+        override val messageCollector: MessageCollector,
+        override val incrementalComponents: IncrementalCompilationComponents?,
+        override val extensionRegistrars: List<FirExtensionRegistrar>,
+        override val configuration: CompilerConfiguration,
+        override val targetIds: List<TargetId>?
+    ) : FrontendContext {
+        constructor(
+            projectEnvironment: AbstractProjectEnvironment,
+            environment: KotlinCoreEnvironment,
+            projectConfiguration: CompilerConfiguration,
+            project: Project?
+        ) : this(
+            projectEnvironment,
+            environment.messageCollector,
+            incrementalComponents = projectConfiguration.get(JVMConfigurationKeys.INCREMENTAL_COMPILATION_COMPONENTS),
+            extensionRegistrars = project?.let { FirExtensionRegistrar.getInstances(it) } ?: emptyList(),
+            configuration = projectConfiguration,
+            targetIds = projectConfiguration.get(JVMConfigurationKeys.MODULES)?.map(::TargetId)
+        )
+    }
+
+    interface FrontendContext {
+        val projectEnvironment: AbstractProjectEnvironment
+        val messageCollector: MessageCollector
+        val incrementalComponents: IncrementalCompilationComponents?
+        val extensionRegistrars: List<FirExtensionRegistrar>
+        val configuration: CompilerConfiguration
+        val targetIds: List<TargetId>?
+    }
 }

 fun findMainClass(fir: List<FirFile>): FqName? {
@@ -13,13 +13,19 @@ import com.intellij.psi.PsiJavaModule
 import com.intellij.psi.search.DelegatingGlobalSearchScope
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analyzer.AnalysisResult
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.jvm.JvmBackendExtension
+import org.jetbrains.kotlin.backend.jvm.JvmGeneratorExtensionsImpl
 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.checkKotlinPackageUsageForPsi
 import org.jetbrains.kotlin.cli.common.config.addKotlinSourceRoot
 import org.jetbrains.kotlin.cli.common.fir.FirDiagnosticsCompilerResultsReporter
 import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
 import org.jetbrains.kotlin.cli.common.toLogger
+import org.jetbrains.kotlin.cli.jvm.compiler.FirKotlinToJvmBytecodeCompiler.convertToIrAndActualizeForJvm
+import org.jetbrains.kotlin.cli.jvm.compiler.FirKotlinToJvmBytecodeCompiler.runFrontend
 import org.jetbrains.kotlin.cli.jvm.config.*
 import org.jetbrains.kotlin.cli.jvm.config.ClassicFrontendSpecificJvmConfigurationKeys.JAVA_CLASSES_TRACKER
 import org.jetbrains.kotlin.codegen.ClassBuilderFactories
@@ -28,11 +34,18 @@ import org.jetbrains.kotlin.codegen.DefaultCodegenFactory
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.*
 import org.jetbrains.kotlin.config.CommonConfigurationKeys.LOOKUP_TRACKER
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
 import org.jetbrains.kotlin.diagnostics.impl.BaseDiagnosticsCollector
+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.pipeline.Fir2IrActualizedResult
 import org.jetbrains.kotlin.ir.backend.jvm.jvmResolveLibraries
+import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmIrMangler
 import org.jetbrains.kotlin.load.kotlin.ModuleVisibilityManager
 import org.jetbrains.kotlin.modules.Module
+import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -69,36 +82,92 @@ object KotlinToJVMBytecodeCompiler {
        }

        val projectConfiguration = environment.configuration
-        if (projectConfiguration.getBoolean(CommonConfigurationKeys.USE_FIR)) {
+        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)) {
+            // K2/PSI: base checks
            val projectEnvironment =
                VfsBasedProjectEnvironment(
                    environment.project,
                    VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
                ) { environment.createPackagePartProvider(it) }
-            return FirKotlinToJvmBytecodeCompiler.compileModulesUsingFrontendIR(
-                projectEnvironment,
-                environment.configuration,
-                environment.messageCollector,
-                environment.getSourceFiles(),
-                buildFile, chunk
+
+            if (!FirKotlinToJvmBytecodeCompiler.checkNotSupportedPlugins(
+                    projectConfiguration, environment.messageCollector
+                )
+            ) return false
+
+            // K2/PSI: single chunk mode fallback
+            if (chunk.size == 1) {
+                return FirKotlinToJvmBytecodeCompiler.compileModulesUsingFrontendIR(
+                    projectEnvironment,
+                    environment.configuration,
+                    environment.messageCollector,
+                    environment.getSourceFiles(),
+                    buildFile, chunk
+                )
+            }
+
+            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)
+            }
+        } 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)
        }
-
-        val result = repeatAnalysisIfNeeded(analyze(environment), environment)
-        if (result == null || !result.shouldGenerateCode) return false
-
-        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
-
-        result.throwIfError()
-
-        val mainClassFqName =
-            if (chunk.size == 1 && projectConfiguration.get(JVMConfigurationKeys.OUTPUT_JAR) != null)
-                findMainClass(result.bindingContext, projectConfiguration.languageVersionSettings, environment.getSourceFiles())
-            else null
-
+        // K1/K2 common multi-chunk part
        val localFileSystem = VirtualFileManager.getInstance().getFileSystem(StandardFileSystems.FILE_PROTOCOL)
+        val (codegenFactory, wholeBackendInput) = irResult

-        val (codegenFactory, wholeBackendInput) = convertToIr(environment, result)
        val diagnosticsReporter = DiagnosticReporterFactory.createReporter()

        val codegenInputs = ArrayList<CodegenFactory.CodegenInput>(chunk.size)
@@ -110,16 +179,23 @@ object KotlinToJVMBytecodeCompiler {
            if (!checkKotlinPackageUsageForPsi(environment.configuration, ktFiles)) return false
            val moduleConfiguration = projectConfiguration.applyModuleProperties(module, buildFile)

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

        val outputs = ArrayList<GenerationState>(chunk.size)

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

        return writeOutputsIfNeeded(environment.project, projectConfiguration, chunk, outputs, mainClassFqName)
@@ -195,8 +271,8 @@ object KotlinToJVMBytecodeCompiler {
        val (codegenFactory, backendInput) = convertToIr(environment, result)
        val diagnosticsReporter = DiagnosticReporterFactory.createReporter()
        val input = runLowerings(
-            environment, environment.configuration, result, environment.getSourceFiles(), null, codegenFactory, backendInput,
-            diagnosticsReporter
+            environment, environment.configuration, result.moduleDescriptor, result.bindingContext,
+            environment.getSourceFiles(), null, codegenFactory, backendInput, diagnosticsReporter
        )
        return runCodegen(input, input.state, codegenFactory, result.bindingContext, diagnosticsReporter, environment.configuration)
    }
@@ -227,6 +303,22 @@ object KotlinToJVMBytecodeCompiler {
        return Pair(codegenFactory, backendInput)
    }

+    private fun Fir2IrActualizedResult.codegenFactoryWithJvmIrBackendInput(
+        configuration: CompilerConfiguration
+    ): Pair<CodegenFactory, CodegenFactory.BackendInput> {
+        val phaseConfig = configuration.get(CLIConfigurationKeys.PHASE_CONFIG)
+        val codegenFactory = JvmIrCodegenFactory(configuration, phaseConfig)
+        return codegenFactory to JvmIrCodegenFactory.JvmIrBackendInput(
+            irModuleFragment,
+            components.symbolTable,
+            phaseConfig,
+            components.irProviders,
+            JvmGeneratorExtensionsImpl(configuration),
+            JvmBackendExtension.Default,
+            pluginContext,
+        ) {}
+    }
+
    fun analyze(environment: KotlinCoreEnvironment): AnalysisResult? {
        val collector = environment.messageCollector
        val sourceFiles = environment.getSourceFiles()
@@ -296,26 +388,32 @@ object KotlinToJVMBytecodeCompiler {
    private fun runLowerings(
        environment: KotlinCoreEnvironment,
        configuration: CompilerConfiguration,
-        result: AnalysisResult,
+        moduleDescriptor: ModuleDescriptor,
+        bindingContext: BindingContext,
        sourceFiles: List<KtFile>,
        module: Module?,
        codegenFactory: CodegenFactory,
        backendInput: CodegenFactory.BackendInput,
        diagnosticsReporter: BaseDiagnosticsCollector,
+        firJvmBackendClassResolver: FirJvmBackendClassResolver? = null,
    ): CodegenFactory.CodegenInput {
        val performanceManager = environment.configuration[CLIConfigurationKeys.PERF_MANAGER]

        val state = GenerationState.Builder(
            environment.project,
            ClassBuilderFactories.BINARIES,
-            result.moduleDescriptor,
-            result.bindingContext,
+            moduleDescriptor,
+            bindingContext,
            configuration
        )
            .withModule(module)
            .onIndependentPartCompilationEnd(createOutputFilesFlushingCallbackIfPossible(configuration))
            .diagnosticReporter(diagnosticsReporter)
-            .build()
+            .apply {
+                if (firJvmBackendClassResolver != null) {
+                    jvmBackendClassResolver(firJvmBackendClassResolver)
+                }
+            }.build()

        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()