FIR JS: Add K2 support to K2JsIrCompiler

2022-08-08 14:16:40 +02:00
parent 4e8fed09b4
commit 7a57e8d903
6 changed files with 353 additions and 120 deletions
@@ -396,7 +396,7 @@ class K2JSCompilerArguments : CommonCompilerArguments() {
 }

 fun K2JSCompilerArguments.isPreIrBackendDisabled(): Boolean =
-    irOnly || irProduceJs || irProduceKlibFile || irBuildCache
+    irOnly || irProduceJs || irProduceKlibFile || irBuildCache || useK2

 fun K2JSCompilerArguments.isIrBackendEnabled(): Boolean =
-    irProduceKlibDir || irProduceJs || irProduceKlibFile || wasm || irBuildCache
+    irProduceKlibDir || irProduceJs || irProduceKlibFile || wasm || irBuildCache || useK2
@@ -172,10 +172,6 @@ public class K2JSCompiler extends CLICompiler<K2JSCompilerArguments> {
            @Nullable KotlinPaths paths
    ) {
        MessageCollector messageCollector = configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY);
-        if (configuration.getBoolean(CommonConfigurationKeys.USE_FIR)) {
-            messageCollector.report(ERROR, "K2 does not support JS target right now", null);
-            return ExitCode.COMPILATION_ERROR;
-        }

        ExitCode exitCode = OK;

@@ -8,12 +8,17 @@ package org.jetbrains.kotlin.cli.js
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.backend.common.CompilationException
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.common.serialization.metadata.KlibMetadataVersion
-import org.jetbrains.kotlin.backend.wasm.compileWasm
+import org.jetbrains.kotlin.backend.common.serialization.signature.IdSignatureDescriptor
 import org.jetbrains.kotlin.backend.wasm.compileToLoweredIr
+import org.jetbrains.kotlin.backend.wasm.compileWasm
+import org.jetbrains.kotlin.backend.wasm.dce.eliminateDeadDeclarations
 import org.jetbrains.kotlin.backend.wasm.wasmPhases
 import org.jetbrains.kotlin.backend.wasm.writeCompilationResult
+import org.jetbrains.kotlin.builtins.DefaultBuiltIns
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.cli.common.*
 import org.jetbrains.kotlin.cli.common.ExitCode.*
 import org.jetbrains.kotlin.cli.common.arguments.K2JSCompilerArguments
@@ -21,43 +26,77 @@ import org.jetbrains.kotlin.cli.common.arguments.K2JsArgumentConstants
 import org.jetbrains.kotlin.cli.common.arguments.K2JsArgumentConstants.RUNTIME_DIAGNOSTIC_EXCEPTION
 import org.jetbrains.kotlin.cli.common.arguments.K2JsArgumentConstants.RUNTIME_DIAGNOSTIC_LOG
 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.messages.CompilerMessageLocation
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity.*
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.cli.common.messages.MessageUtil
+import org.jetbrains.kotlin.cli.js.klib.generateIrForKlibSerialization
 import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.cli.jvm.plugins.PluginCliParser
+import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.Services
+import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
+import org.jetbrains.kotlin.fir.DependencyListForCliModule
+import org.jetbrains.kotlin.fir.FirModuleDataImpl
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.backend.Fir2IrConverter
+import org.jetbrains.kotlin.fir.backend.Fir2IrExtensions
+import org.jetbrains.kotlin.fir.backend.Fir2IrVisibilityConverter
+import org.jetbrains.kotlin.fir.backend.jvm.Fir2IrJvmSpecialAnnotationSymbolProvider
+import org.jetbrains.kotlin.fir.backend.jvm.FirJvmKotlinMangler
+import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
+import org.jetbrains.kotlin.fir.descriptors.FirModuleDescriptor
+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.pipeline.buildFirFromKtFiles
+import org.jetbrains.kotlin.fir.pipeline.runCheckers
+import org.jetbrains.kotlin.fir.pipeline.runResolution
+import org.jetbrains.kotlin.fir.resolve.providers.firProvider
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirProviderImpl
+import org.jetbrains.kotlin.fir.session.FirSessionConfigurator
+import org.jetbrains.kotlin.fir.session.FirJsSessionFactory
 import org.jetbrains.kotlin.incremental.components.ExpectActualTracker
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.incremental.js.IncrementalDataProvider
 import org.jetbrains.kotlin.incremental.js.IncrementalNextRoundChecker
 import org.jetbrains.kotlin.incremental.js.IncrementalResultsConsumer
-import org.jetbrains.kotlin.backend.wasm.dce.eliminateDeadDeclarations
-import org.jetbrains.kotlin.cli.js.klib.generateIrForKlibSerialization
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.ir.backend.js.*
 import org.jetbrains.kotlin.ir.backend.js.codegen.JsGenerationGranularity
 import org.jetbrains.kotlin.ir.backend.js.ic.*
+import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsManglerDesc
+import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsManglerIr
 import org.jetbrains.kotlin.ir.backend.js.transformers.irToJs.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImplForJsIC
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
+import org.jetbrains.kotlin.ir.util.IrMessageLogger
 import org.jetbrains.kotlin.js.analyzer.JsAnalysisResult
 import org.jetbrains.kotlin.js.config.*
+import org.jetbrains.kotlin.js.resolve.JsPlatformAnalyzerServices
+import org.jetbrains.kotlin.library.KotlinAbiVersion
+import org.jetbrains.kotlin.library.unresolvedDependencies
 import org.jetbrains.kotlin.metadata.deserialization.BinaryVersion
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.js.JsPlatforms
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.serialization.js.ModuleKind
+import org.jetbrains.kotlin.storage.LockBasedStorageManager
 import org.jetbrains.kotlin.utils.KotlinPaths
 import org.jetbrains.kotlin.utils.PathUtil
 import org.jetbrains.kotlin.utils.join
 import java.io.File
 import java.io.IOException
+import java.nio.file.Paths

 private val K2JSCompilerArguments.granularity: JsGenerationGranularity
    get() = when {
@@ -129,10 +168,6 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
        paths: KotlinPaths?
    ): ExitCode {
        val messageCollector = configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY)
-        if (configuration.getBoolean(CommonConfigurationKeys.USE_FIR)) {
-            messageCollector.report(ERROR, "K2 does not support JS target right now")
-            return COMPILATION_ERROR
-        }

        val pluginLoadResult = loadPlugins(paths, arguments, configuration)
        if (pluginLoadResult != OK) return pluginLoadResult
@@ -234,116 +269,28 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
        // TODO: Handle non-empty main call arguments
        val mainCallArguments = if (K2JsArgumentConstants.NO_CALL == arguments.main) null else emptyList<String>()

-        val cacheDirectories = configureLibraries(arguments.cacheDirectories)
-
-        // TODO: Use JS IR IC infrastructure for WASM?
-        val icCaches = if (!arguments.wasm && cacheDirectories.isNotEmpty()) {
-            messageCollector.report(INFO, "")
-            messageCollector.report(INFO, "Building cache:")
-            messageCollector.report(INFO, "to: $outputDir")
-            messageCollector.report(INFO, arguments.cacheDirectories ?: "")
-            messageCollector.report(INFO, libraries.toString())
-
-            val start = System.currentTimeMillis()
-
-            val cacheUpdater = CacheUpdater(
-                mainModule = arguments.includes!!,
-                allModules = libraries,
-                icCachePaths = cacheDirectories,
-                compilerConfiguration = configurationJs,
-                irFactory = { IrFactoryImplForJsIC(WholeWorldStageController()) },
-                mainArguments = mainCallArguments,
-                compilerInterfaceFactory = { mainModule, cfg -> JsIrCompilerWithIC(mainModule, cfg, arguments.granularity) }
-            )
-
-            val artifacts = cacheUpdater.actualizeCaches()
-            messageCollector.report(INFO, "IC rebuilt overall time: ${System.currentTimeMillis() - start}ms")
-            for ((event, duration) in cacheUpdater.getStopwatchLaps()) {
-                messageCollector.report(INFO, "  $event: ${(duration / 1e6).toInt()}ms")
-            }
-
-            var libIndex = 0
-            for ((libFile, srcFiles) in cacheUpdater.getDirtyFileStats()) {
-                val (msg, showFiles) = when {
-                    srcFiles.values.all { it.contains(DirtyFileState.ADDED_FILE) } -> "fully rebuilt due to clean build" to false
-                    srcFiles.values.all { it.contains(DirtyFileState.MODIFIED_CONFIG) } -> "fully rebuilt due to config modification" to false
-                    else -> "partially rebuilt" to true
-                }
-                messageCollector.report(INFO, "${++libIndex}) module [${File(libFile.path).name}] was $msg")
-                if (showFiles) {
-                    var fileIndex = 0
-                    for ((srcFile, stat) in srcFiles) {
-                        val statStr = stat.joinToString { it.str }
-                        // Use index, because MessageCollector ignores already reported messages
-                        messageCollector.report(INFO, "  $libIndex.${++fileIndex}) file [${File(srcFile.path).name}]: ($statStr)")
-                    }
-                }
-            }
-
-            artifacts
-        } else emptyList()
+        val icCaches = prepareIcCaches(arguments, messageCollector, outputDir, libraries, configurationJs, mainCallArguments)

        // Run analysis if main module is sources
-        lateinit var sourceModule: ModulesStructure
+        var sourceModule: ModulesStructure? = null
        val includes = arguments.includes
        if (includes == null) {
-            do {
-                sourceModule = prepareAnalyzedSourceModule(
-                    projectJs,
-                    environmentForJS.getSourceFiles(),
-                    configurationJs,
-                    libraries,
-                    friendLibraries,
-                    AnalyzerWithCompilerReport(configurationJs)
+            val outputKlibPath =
+                if (arguments.irProduceKlibFile) outputDir.resolve("$outputName.klib").normalize().absolutePath
+                else outputDirPath
+            if (arguments.useK2) {
+                messageCollector.report(
+                    STRONG_WARNING,
+                    "ATTENTION!\n This build uses experimental K2 compiler: \n  -Xuse-k2"
                )
-                val result = sourceModule.jsFrontEndResult.jsAnalysisResult
-                if (result is JsAnalysisResult.RetryWithAdditionalRoots) {
-                    environmentForJS.addKotlinSourceRoots(result.additionalKotlinRoots)
-                }
-            } while (result is JsAnalysisResult.RetryWithAdditionalRoots)
-            if (!sourceModule.jsFrontEndResult.jsAnalysisResult.shouldGenerateCode)
-                return OK
-        }
+                sourceModule = processSourceModuleWithK2(environmentForJS, libraries, friendLibraries, arguments, outputKlibPath)
+            } else {
+                sourceModule = processSourceModule(environmentForJS, libraries, friendLibraries, arguments, outputKlibPath)

-        if (arguments.irProduceKlibDir || arguments.irProduceKlibFile) {
-            if (arguments.irProduceKlibFile) {
-                require(outputFile.extension == KLIB_FILE_EXTENSION) { "Please set up .klib file as output" }
-            }
-            val moduleSourceFiles = (sourceModule.mainModule as MainModule.SourceFiles).files
-            val icData = configurationJs.incrementalDataProvider?.getSerializedData(moduleSourceFiles) ?: emptyList()
-            val expectDescriptorToSymbol = mutableMapOf<DeclarationDescriptor, IrSymbol>()
-
-            val moduleFragment = generateIrForKlibSerialization(
-                projectJs,
-                moduleSourceFiles,
-                configurationJs,
-                sourceModule.jsFrontEndResult.jsAnalysisResult,
-                sortDependencies(sourceModule.descriptors),
-                icData,
-                expectDescriptorToSymbol,
-                IrFactoryImpl,
-                verifySignatures = true
-            ) {
-                sourceModule.getModuleDescriptor(it)
+                if (!sourceModule.jsFrontEndResult.jsAnalysisResult.shouldGenerateCode)
+                    return OK
            }

-            val metadataSerializer =
-                KlibMetadataIncrementalSerializer(configuration, sourceModule.project, sourceModule.jsFrontEndResult.hasErrors)
-
-            generateKLib(
-                sourceModule,
-                outputKlibPath = if (arguments.irProduceKlibFile)
-                    outputDir.resolve("$outputName.klib").normalize().absolutePath
-                else
-                    outputDirPath,
-                nopack = arguments.irProduceKlibDir,
-                jsOutputName = arguments.irPerModuleOutputName,
-                icData = icData,
-                expectDescriptorToSymbol = expectDescriptorToSymbol,
-                moduleFragment = moduleFragment
-            ) { file ->
-                metadataSerializer.serializeScope(file, sourceModule.jsFrontEndResult.bindingContext, moduleFragment.descriptor)
-            }
        }

        if (arguments.irProduceJs) {
@@ -394,7 +341,7 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                    friendLibraries
                )
            } else {
-                sourceModule
+                sourceModule!!
            }

            if (arguments.wasm) {
@@ -457,6 +404,288 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
        return OK
    }

+    private fun processSourceModule(
+        environmentForJS: KotlinCoreEnvironment,
+        libraries: List<String>,
+        friendLibraries: List<String>,
+        arguments: K2JSCompilerArguments,
+        outputKlibPath: String
+    ): ModulesStructure {
+        lateinit var sourceModule: ModulesStructure
+        do {
+            sourceModule = prepareAnalyzedSourceModule(
+                environmentForJS.project,
+                environmentForJS.getSourceFiles(),
+                environmentForJS.configuration,
+                libraries,
+                friendLibraries,
+                AnalyzerWithCompilerReport(environmentForJS.configuration)
+            )
+            val result = sourceModule.jsFrontEndResult.jsAnalysisResult
+            if (result is JsAnalysisResult.RetryWithAdditionalRoots) {
+                environmentForJS.addKotlinSourceRoots(result.additionalKotlinRoots)
+            }
+        } while (result is JsAnalysisResult.RetryWithAdditionalRoots)
+
+        if (sourceModule.jsFrontEndResult.jsAnalysisResult.shouldGenerateCode && (arguments.irProduceKlibDir || arguments.irProduceKlibFile)) {
+            val moduleSourceFiles = (sourceModule.mainModule as MainModule.SourceFiles).files
+            val icData = environmentForJS.configuration.incrementalDataProvider?.getSerializedData(moduleSourceFiles) ?: emptyList()
+            val expectDescriptorToSymbol = mutableMapOf<DeclarationDescriptor, IrSymbol>()
+
+            val moduleFragment = generateIrForKlibSerialization(
+                environmentForJS.project,
+                moduleSourceFiles,
+                environmentForJS.configuration,
+                sourceModule.jsFrontEndResult.jsAnalysisResult,
+                sortDependencies(sourceModule.descriptors),
+                icData,
+                expectDescriptorToSymbol,
+                IrFactoryImpl,
+                verifySignatures = true
+            ) {
+                sourceModule.getModuleDescriptor(it)
+            }
+
+            val metadataSerializer =
+                KlibMetadataIncrementalSerializer(environmentForJS.configuration, sourceModule.project, sourceModule.jsFrontEndResult.hasErrors)
+
+            generateKLib(
+                sourceModule,
+                outputKlibPath,
+                nopack = arguments.irProduceKlibDir,
+                jsOutputName = arguments.irPerModuleOutputName,
+                icData = icData,
+                expectDescriptorToSymbol = expectDescriptorToSymbol,
+                moduleFragment = moduleFragment
+            ) { file ->
+                metadataSerializer.serializeScope(file, sourceModule.jsFrontEndResult.bindingContext, moduleFragment.descriptor)
+            }
+        }
+        return sourceModule
+    }
+
+    private fun processSourceModuleWithK2(
+        environmentForJS: KotlinCoreEnvironment,
+        libraries: List<String>,
+        friendLibraries: List<String>,
+        arguments: K2JSCompilerArguments,
+        outputKlibPath: String
+    ): ModulesStructure? {
+        val configuration = environmentForJS.configuration
+        val messageCollector = configuration.getNotNull(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY)
+        val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter()
+        val renderDiagnosticNames = configuration.getBoolean(CLIConfigurationKeys.RENDER_DIAGNOSTIC_INTERNAL_NAME)
+
+        // FIR
+
+        val sessionProvider = FirProjectSessionProvider()
+        val extensionRegistrars = FirExtensionRegistrar.getInstances(environmentForJS.project)
+        val sessionConfigurator: FirSessionConfigurator.() -> Unit = {
+            if (arguments.extendedCompilerChecks) {
+                registerExtendedCommonCheckers()
+            }
+        }
+
+        val mainModuleName = Name.special("<${configuration.get(CommonConfigurationKeys.MODULE_NAME)!!}>")
+
+        val ktFiles = environmentForJS.getSourceFiles()
+        val syntaxErrors = ktFiles.fold(false) { errorsFound, ktFile ->
+            AnalyzerWithCompilerReport.reportSyntaxErrors(ktFile, messageCollector).isHasErrors or errorsFound
+        }
+
+        val mainModule = MainModule.SourceFiles(environmentForJS.getSourceFiles())
+
+        val dependencyList = DependencyListForCliModule.build(mainModuleName, JsPlatforms.defaultJsPlatform, JsPlatformAnalyzerServices) {
+            dependencies(libraries.map { Paths.get(it).toAbsolutePath() })
+            friendDependencies(friendLibraries.map { Paths.get(it).toAbsolutePath() })
+            // TODO: !!! dependencies module data?
+        }
+
+        val repositories = configuration[JSConfigurationKeys.REPOSITORIES] ?: emptyList()
+        val logger = configuration.resolverLogger
+        val resolvedLibraries = jsResolveLibraries(libraries + friendLibraries, repositories, logger).getFullResolvedList()
+
+        FirJsSessionFactory.createJsLibrarySession(
+            mainModuleName,
+            resolvedLibraries,
+            sessionProvider,
+            dependencyList.moduleDataProvider,
+            configuration.languageVersionSettings
+        )
+
+        val mainModuleData = FirModuleDataImpl(
+            mainModuleName,
+            dependencyList.regularDependencies,
+            dependencyList.dependsOnDependencies,
+            dependencyList.friendsDependencies,
+            dependencyList.platform,
+            dependencyList.analyzerServices
+        )
+
+        val session = FirJsSessionFactory.createJsModuleBasedSession(
+            mainModuleData,
+            sessionProvider,
+            extensionRegistrars,
+            configuration.languageVersionSettings,
+            null,
+            sessionConfigurator,
+        )
+
+        val rawFirFiles = session.buildFirFromKtFiles(ktFiles)
+        val (scopeSession, firFiles) = session.runResolution(rawFirFiles)
+        session.runCheckers(scopeSession, firFiles, diagnosticsReporter)
+
+        if (syntaxErrors || diagnosticsReporter.hasErrors) {
+            FirDiagnosticsCompilerResultsReporter.reportToMessageCollector(diagnosticsReporter, messageCollector, renderDiagnosticNames)
+            return null
+        }
+
+        // FIR2IR
+
+        val fir2IrExtensions = Fir2IrExtensions.Default
+        val signaturer = IdSignatureDescriptor(JsManglerDesc)
+        val commonFirFiles = session.moduleData.dependsOnDependencies
+            .map { it.session }
+            .filter { it.kind == FirSession.Kind.Source }
+            .flatMap { (it.firProvider as FirProviderImpl).getAllFirFiles() }
+
+        var builtInsModule: KotlinBuiltIns? = null
+        val dependencies = mutableListOf<ModuleDescriptorImpl>()
+
+        val librariesDescriptors = resolvedLibraries.map { resolvedLibrary ->
+            val storageManager = LockBasedStorageManager("ModulesStructure")
+
+            val moduleDescriptor = JsFactories.DefaultDeserializedDescriptorFactory.createDescriptorOptionalBuiltIns(
+                resolvedLibrary.library,
+                configuration.languageVersionSettings,
+                storageManager,
+                builtInsModule,
+                packageAccessHandler = null,
+                lookupTracker = LookupTracker.DO_NOTHING
+            )
+            dependencies += moduleDescriptor
+            moduleDescriptor.setDependencies(ArrayList(dependencies))
+
+            val isBuiltIns = resolvedLibrary.library.unresolvedDependencies.isEmpty()
+            if (isBuiltIns) builtInsModule = moduleDescriptor.builtIns
+
+            moduleDescriptor
+        }
+
+        val fir2irResult = Fir2IrConverter.createModuleFragmentWithSignaturesIfNeeded(
+            session, scopeSession, firFiles + commonFirFiles,
+            configuration.languageVersionSettings, signaturer,
+            fir2IrExtensions,
+            FirJvmKotlinMangler(session), // TODO: replace with potentially simpler JS version
+            JsManglerIr, IrFactoryImpl,
+            Fir2IrVisibilityConverter.Default,
+            Fir2IrJvmSpecialAnnotationSymbolProvider(), // TODO: replace with appropriate (probably empty) implementation
+            IrGenerationExtension.getInstances(environmentForJS.project),
+            generateSignatures = false,
+            kotlinBuiltIns = builtInsModule ?: DefaultBuiltIns.Instance // TODO: consider passing externally
+        ).also {
+            (it.irModuleFragment.descriptor as? FirModuleDescriptor)?.let { it.allDependencyModules = librariesDescriptors }
+        }
+
+        // Serialize klib
+
+        if (/*sourceModule.jsFrontEndResult.jsAnalysisResult.shouldGenerateCode && */(arguments.irProduceKlibDir || arguments.irProduceKlibFile)) {
+            val sourceFiles = firFiles.mapNotNull { it.sourceFile }
+            val firFilesBySourceFile = firFiles.associateBy { it.sourceFile }
+
+            val icData = environmentForJS.configuration.incrementalDataProvider?.getSerializedData(sourceFiles) ?: emptyList()
+            // TODO: expect -> actual mapping
+            val expectDescriptorToSymbol = mutableMapOf<DeclarationDescriptor, IrSymbol>()
+
+            val metadataVersion =
+                configuration.get(CommonConfigurationKeys.METADATA_VERSION)
+                    ?: GenerationState.LANGUAGE_TO_METADATA_VERSION.getValue(configuration.languageVersionSettings.languageVersion)
+
+            serializeModuleIntoKlib(
+                configuration[CommonConfigurationKeys.MODULE_NAME]!!,
+                configuration,
+                configuration.get(IrMessageLogger.IR_MESSAGE_LOGGER) ?: IrMessageLogger.None,
+                sourceFiles,
+                klibPath = outputKlibPath,
+                resolvedLibraries.map { it.library },
+                fir2irResult.irModuleFragment,
+                expectDescriptorToSymbol,
+                cleanFiles = icData,
+                nopack = true,
+                perFile = false,
+                containsErrorCode = messageCollector.hasErrors() || diagnosticsReporter.hasErrors,
+                abiVersion = KotlinAbiVersion.CURRENT, // TODO get from test file data
+                jsOutputName = null
+            ) { file ->
+                val firFile = firFilesBySourceFile[file] ?: error("cannot find FIR file by source file ${file.name} (${file.path})")
+                serializeSingleFirFile(firFile, session, scopeSession, metadataVersion)
+            }
+        }
+
+        return ModulesStructure(
+            environmentForJS.project, mainModule, configuration, libraries, friendLibraries
+        )
+    }
+
+    private fun prepareIcCaches(
+        arguments: K2JSCompilerArguments,
+        messageCollector: MessageCollector,
+        outputDir: File,
+        libraries: List<String>,
+        configurationJs: CompilerConfiguration,
+        mainCallArguments: List<String>?
+    ): List<ModuleArtifact> {
+        val cacheDirectories = configureLibraries(arguments.cacheDirectories)
+
+        // TODO: Use JS IR IC infrastructure for WASM?
+        val icCaches = if (!arguments.wasm && cacheDirectories.isNotEmpty()) {
+            messageCollector.report(INFO, "")
+            messageCollector.report(INFO, "Building cache:")
+            messageCollector.report(INFO, "to: ${outputDir}")
+            messageCollector.report(INFO, arguments.cacheDirectories ?: "")
+            messageCollector.report(INFO, libraries.toString())
+
+            val start = System.currentTimeMillis()
+
+            val cacheUpdater = CacheUpdater(
+                mainModule = arguments.includes!!,
+                allModules = libraries,
+                icCachePaths = cacheDirectories,
+                compilerConfiguration = configurationJs,
+                irFactory = { IrFactoryImplForJsIC(WholeWorldStageController()) },
+                mainArguments = mainCallArguments,
+                compilerInterfaceFactory = { mainModule, cfg -> JsIrCompilerWithIC(mainModule, cfg, arguments.granularity) }
+            )
+
+            val artifacts = cacheUpdater.actualizeCaches()
+            messageCollector.report(INFO, "IC rebuilt overall time: ${System.currentTimeMillis() - start}ms")
+            for ((event, duration) in cacheUpdater.getStopwatchLaps()) {
+                messageCollector.report(INFO, "  $event: ${(duration / 1e6).toInt()}ms")
+            }
+
+            var libIndex = 0
+            for ((libFile, srcFiles) in cacheUpdater.getDirtyFileStats()) {
+                val (msg, showFiles) = when {
+                    srcFiles.values.all { it.contains(DirtyFileState.ADDED_FILE) } -> "fully rebuilt due to clean build" to false
+                    srcFiles.values.all { it.contains(DirtyFileState.MODIFIED_CONFIG) } -> "fully rebuilt due to config modification" to false
+                    else -> "partially rebuilt" to true
+                }
+                messageCollector.report(INFO, "${++libIndex}) module [${File(libFile.path).name}] was $msg")
+                if (showFiles) {
+                    var fileIndex = 0
+                    for ((srcFile, stat) in srcFiles) {
+                        val statStr = stat.joinToString { it.str }
+                        // Use index, because MessageCollector ignores already reported messages
+                        messageCollector.report(INFO, "  $libIndex.${++fileIndex}) file [${File(srcFile.path).name}]: ($statStr)")
+                    }
+                }
+            }
+
+            artifacts
+        } else emptyList()
+        return icCaches
+    }
+
    private fun CompilationOutputs.write(outputDir: File, outputName: String) {
        val outputFile = outputDir.resolve("$outputName.js")
        outputFile.parentFile.mkdirs()
@@ -1,4 +1,6 @@
 $TESTDATA_DIR$/fir.kt
 -Xuse-k2
-output
+-ir-output-dir
 $TEMP_DIR$
+-ir-output-name
+fir
@@ -1,2 +1,4 @@
-error: k2 does not support JS target right now
-COMPILATION_ERROR
+warning: ATTENTION!
+ This build uses experimental K2 compiler:
+  -Xuse-k2
+OK
@@ -44,6 +44,10 @@ internal open class BaseKotlin2JsCompileConfig<TASK : Kotlin2JsCompile>(
                }
            )

+            if (propertiesProvider.useK2 == true) {
+                task.kotlinOptions.useK2 = true
+            }
+
            task.destinationDirectory
                .convention(
                    project.objects.directoryProperty().fileProvider(