[Wasm] Loader improvements

- Output ES modules instead of plain files
- Support -Xwasm-launcher=d8 for d8 shell used in tests and benchmarks.
- Reuse launcher generation logic in CLI and box tests runners.
- Create separate output directory for each box since
  there are multiple output files generated for each test.
- Stop using absolute paths in generate JS files
  to simplify running generated code on different machine
- Remove ">>>" from println output


Merge-request: KT-MR-5729
Merged-by: Svyatoslav Kuzmich <svyatoslav.kuzmich@jetbrains.com>

compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JSCompilerArguments.kt

@@ -242,7 +242,7 @@ class K2JSCompilerArguments : CommonCompilerArguments() {
 
     @Argument(
             value = "-Xwasm-launcher",
-            valueDescription = "esm|nodejs",
+            valueDescription = "esm|nodejs|d8",
             description = "Picks flavor for the wasm launcher. Default is ESM."
     )
     var wasmLauncher: String? by NullableStringFreezableVar("esm")

compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt

@@ -11,9 +11,11 @@ import com.intellij.openapi.util.text.StringUtil
 import org.jetbrains.kotlin.backend.common.CompilationException
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.common.serialization.metadata.KlibMetadataVersion
+import org.jetbrains.kotlin.backend.wasm.WasmLoaderKind
 import org.jetbrains.kotlin.backend.wasm.compileWasm
 import org.jetbrains.kotlin.backend.wasm.compileToLoweredIr
 import org.jetbrains.kotlin.backend.wasm.wasmPhases
+import org.jetbrains.kotlin.backend.wasm.writeCompilationResult
 import org.jetbrains.kotlin.cli.common.*
 import org.jetbrains.kotlin.cli.common.ExitCode.*
 import org.jetbrains.kotlin.cli.common.arguments.K2JSCompilerArguments
@@ -38,6 +40,7 @@ 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.ir.backend.js.*
 import org.jetbrains.kotlin.ir.backend.js.codegen.JsGenerationGranularity
 import org.jetbrains.kotlin.ir.backend.js.ic.actualizeCaches
@@ -326,49 +329,30 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                     exportedDeclarations = setOf(FqName("main")),
                     propertyLazyInitialization = arguments.irPropertyLazyInitialization,
                 )
+                if (arguments.irDce) {
+                    eliminateDeadDeclarations(allModules, backendContext)
+                }
                 val res = compileWasm(
                     allModules = allModules,
                     backendContext = backendContext,
                     emitNameSection = arguments.wasmDebug,
-                    dceEnabled = arguments.irDce,
+                    allowIncompleteImplementations = arguments.irDce,
                 )
-                val outputWasmFile = outputFile.withReplacedExtensionOrNull(outputFile.extension, "wasm")!!
-                outputWasmFile.writeBytes(res.wasm)
-                val outputWatFile = outputFile.withReplacedExtensionOrNull(outputFile.extension, "wat")!!
-                outputWatFile.writeText(res.wat)
 
-                val esmRunner = """
-                    export default WebAssembly.instantiateStreaming(fetch('${outputWasmFile.name}'), { runtime, js_code }).then((it) => {
-                        wasmInstance = it.instance;
-                        wasmInstance.exports.__init?.();
-                        wasmInstance.exports.startUnitTests?.();
-                        
-                        return it.instance.exports;
-                    });
-                """.trimIndent()
-
-                val nodeRunner = """
-                    const fs = require('fs');
-                    var path = require('path');
-                    const wasmBuffer = fs.readFileSync(path.resolve(__dirname, './${outputWasmFile.name}'));
-                    
-                    module.exports = WebAssembly.instantiate(wasmBuffer, { runtime, js_code }).then(wasm => {
-                        wasmInstance = wasm.instance;
-                    
-                        wasmInstance.exports.__init?.();
-                        wasmInstance.exports.startUnitTests?.();
-                    
-                        return wasmInstance.exports
-                    });
-                """.trimIndent()
-
-                val runner = when (arguments.wasmLauncher) {
-                    "esm" -> esmRunner
-                    "nodejs" -> nodeRunner
+                val launcherKind = when (arguments.wasmLauncher) {
+                    "esm" -> WasmLoaderKind.BROWSER
+                    "nodejs" -> WasmLoaderKind.NODE
+                    "d8" -> WasmLoaderKind.D8
                     else -> throw IllegalArgumentException("Unrecognized flavor for the wasm launcher")
                 }
 
-                outputFile.writeText(res.js + "\n" + runner)
+                writeCompilationResult(
+                    result = res,
+                    dir = outputFile.parentFile,
+                    loaderKind = launcherKind,
+                    fileNameBase = outputFile.nameWithoutExtension
+                )
+
                 return OK
             }
 

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/compiler.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.backend.wasm
 
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
-import org.jetbrains.kotlin.backend.wasm.dce.eliminateDeadDeclarations
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.WasmCompiledModuleFragment
 import org.jetbrains.kotlin.backend.wasm.ir2wasm.WasmModuleFragmentGenerator
 import org.jetbrains.kotlin.backend.wasm.lower.markExportedDeclarations
@@ -23,6 +22,7 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.wasm.ir.convertors.WasmIrToBinary
 import org.jetbrains.kotlin.wasm.ir.convertors.WasmIrToText
 import java.io.ByteArrayOutputStream
+import java.io.File
 
 class WasmCompilerResult(val wat: String, val js: String, val wasm: ByteArray)
 
@@ -75,15 +75,10 @@ fun compileWasm(
     allModules: List<IrModuleFragment>,
     backendContext: WasmBackendContext,
     emitNameSection: Boolean = false,
-    dceEnabled: Boolean = false,
+    allowIncompleteImplementations: Boolean = false,
 ): WasmCompilerResult {
-
-    if (dceEnabled) {
-        eliminateDeadDeclarations(allModules, backendContext)
-    }
-
     val compiledWasmModule = WasmCompiledModuleFragment(backendContext.irBuiltIns)
-    val codeGenerator = WasmModuleFragmentGenerator(backendContext, compiledWasmModule, allowIncompleteImplementations = dceEnabled)
+    val codeGenerator = WasmModuleFragmentGenerator(backendContext, compiledWasmModule, allowIncompleteImplementations = allowIncompleteImplementations)
     allModules.forEach { codeGenerator.generateModule(it) }
 
     val linkedModule = compiledWasmModule.linkWasmCompiledFragments()
@@ -107,13 +102,76 @@ fun compileWasm(
 fun WasmCompiledModuleFragment.generateJs(): String {
     //language=js
     val runtime = """
-    var wasmInstance = null;
     
     const externrefBoxes = new WeakMap();
     """.trimIndent()
 
-    val jsFuns = jsFuns.joinToString(",\n") { "\"" + it.importName + "\" : " + it.jsCode }
-    val jsCode = "\nconst js_code = {$jsFuns};"
+    val jsCodeBody = jsFuns.joinToString(",\n") { "\"" + it.importName + "\" : " + it.jsCode }
+    val jsCodeBodyIndented = jsCodeBody.prependIndent("    ")
+    val jsCode =
+        "\nconst js_code = {\n$jsCodeBodyIndented\n};\n"
 
     return runtime + jsCode
 }
+
+enum class WasmLoaderKind {
+    D8,
+    NODE,
+    BROWSER,
+}
+
+fun generateJsWasmLoader(kind: WasmLoaderKind, wasmFilePath: String, externalJs: String): String {
+    val instantiation = when (kind) {
+        WasmLoaderKind.D8 ->
+            """
+                const wasmModule = new WebAssembly.Module(read('$wasmFilePath', 'binary'));
+                const wasmInstance = new WebAssembly.Instance(wasmModule, { js_code });
+            """.trimIndent()
+
+        WasmLoaderKind.NODE ->
+            """
+                const fs = require('fs');
+                var path = require('path');
+                const wasmBuffer = fs.readFileSync(path.resolve(__dirname, './$wasmFilePath'));
+                const wasmModule = new WebAssembly.Module(wasmBuffer);
+                const wasmInstance = new WebAssembly.Instance(wasmModule, { js_code });
+            """.trimIndent()
+
+        WasmLoaderKind.BROWSER ->
+            """
+                const { wasmInstance } = await WebAssembly.instantiateStreaming(fetch("$wasmFilePath"), { js_code });
+            """.trimIndent()
+    }
+
+    val init =
+        """
+            
+            const wasmExports = wasmInstance.exports;
+            wasmExports.__init();
+            wasmExports.startUnitTests?.();
+            
+        """.trimIndent()
+
+    val export = when (kind) {
+        WasmLoaderKind.D8, WasmLoaderKind.BROWSER ->
+            "export default wasmExports;\n"
+
+        WasmLoaderKind.NODE ->
+            "module.exports = wasmExports;\n"
+    }
+
+    return externalJs + instantiation + init + export
+}
+
+fun writeCompilationResult(
+    result: WasmCompilerResult,
+    dir: File,
+    loaderKind: WasmLoaderKind,
+    fileNameBase: String = "index",
+) {
+    dir.mkdirs()
+    File(dir, "$fileNameBase.wat").writeText(result.wat)
+    File(dir, "$fileNameBase.wasm").writeBytes(result.wasm)
+    val jsWithLoader = generateJsWasmLoader(loaderKind, "./$fileNameBase.wasm", result.js)
+    File(dir, "$fileNameBase.js").writeText(jsWithLoader)
+}

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/dce/Dce.kt

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 import org.jetbrains.kotlin.ir.visitors.acceptVoid
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 
-internal fun eliminateDeadDeclarations(modules: List<IrModuleFragment>, context: WasmBackendContext) {
+fun eliminateDeadDeclarations(modules: List<IrModuleFragment>, context: WasmBackendContext) {
     val printReachabilityInfo =
         context.configuration.getBoolean(JSConfigurationKeys.PRINT_REACHABILITY_INFO) ||
                 java.lang.Boolean.getBoolean("kotlin.wasm.dce.print.reachability.info")

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/lower/JsInteropFunctionsLowering.kt

@@ -364,7 +364,7 @@ class JsInteropFunctionsLowering(val context: WasmBackendContext) : DeclarationT
         val jsCode = buildString {
             append("(f) => (")
             appendParameterList(arity)
-            append(") => wasmInstance.exports.__callFunction_")
+            append(") => wasmExports.__callFunction_")
             append(info.hashString)
             append("(f, ")
             appendParameterList(arity)

compiler/testData/cli/js/jsExtraHelp.out

@@ -39,7 +39,7 @@ where advanced options include:
   -Xwasm                     Use experimental WebAssembly compiler backend
   -Xwasm-debug-info          Add debug info to WebAssembly compiled module
   -Xwasm-kclass-fqn          Enable support for FQ names in KClass
-  -Xwasm-launcher=esm|nodejs Picks flavor for the wasm launcher. Default is ESM.
+  -Xwasm-launcher=esm|nodejs|d8 Picks flavor for the wasm launcher. Default is ESM.
   -Xallow-kotlin-package     Allow compiling code in package 'kotlin' and allow not requiring kotlin.stdlib in module-info
   -Xallow-result-return-type Allow compiling code when `kotlin.Result` is used as a return type
   -Xbuiltins-from-sources    Compile builtIns from sources

compiler/testData/codegen/boxWasmJsInterop/functionTypes.kt

@@ -241,6 +241,14 @@ fun box(): String {
     return "OK"
 }
 
+// TODO: Rewrite test to use module system
+@JsFun("() => { globalThis.main = wasmExports; }")
+external fun hackNonModuleExport()
+
+fun main() {
+    hackNonModuleExport()
+}
+
 // FILE: functionTypes__after.js
 
 const exportedFres = main.exportedF()(1, 20, 300)("<", ">");

compiler/testData/codegen/boxWasmJsInterop/jsExport.kt

@@ -1,3 +1,4 @@
+// IGNORE_BACKEND: JS_IR, JS
 // MODULE: main
 // FILE: externals.kt
 
@@ -25,6 +26,15 @@ fun anyAsEI(any: Any): EI = any as EI
 
 fun box(): String = "OK"
 
+// TODO: Rewrite test to use module system
+
+@JsFun("() => { globalThis.main = wasmExports; }")
+external fun hackNonModuleExport()
+
+fun main() {
+    hackNonModuleExport()
+}
+
 // FILE: jsExport__after.js
 
 const c = main.makeC(300);