From 43eeaa925f373f526eeb90ae4c260c4bbd47e796 Mon Sep 17 00:00:00 2001
From: Igor Laevsky <igor.laevsky@jetbrains.com>
Date: Fri, 22 Oct 2021 20:07:42 +0300
Subject: [PATCH] [Wasm] Add ability to choose flavor for the wasm launcher

---
 .../common/arguments/K2JSCompilerArguments.kt |  7 ++++++
 .../jetbrains/kotlin/cli/js/K2JsIrCompiler.kt | 24 ++++++++++++++++++-
 compiler/testData/cli/js/jsExtraHelp.out      |  1 +
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JSCompilerArguments.kt b/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JSCompilerArguments.kt
index 05465d3fd8c..32509a47e69 100644
--- a/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JSCompilerArguments.kt
+++ b/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JSCompilerArguments.kt
@@ -235,6 +235,13 @@ class K2JSCompilerArguments : CommonCompilerArguments() {
     @Argument(value = "-Xwasm-debug-info", description = "Add debug info to WebAssembly compiled module")
     var wasmDebug: Boolean by FreezableVar(false)
 
+    @Argument(
+            value = "-Xwasm-launcher",
+            valueDescription = "esm|nodejs",
+            description = "Picks flavor for the wasm launcher. Default is ESM."
+    )
+    var wasmLauncher: String? by NullableStringFreezableVar("esm")
+
     override fun configureLanguageFeatures(collector: MessageCollector): MutableMap<LanguageFeature, LanguageFeature.State> {
         return super.configureLanguageFeatures(collector).apply {
             if (extensionFunctionsInExternals) {
diff --git a/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt b/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt
index 6c39c845db4..db5e8c35556 100644
--- a/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt
+++ b/compiler/cli/cli-js/src/org/jetbrains/kotlin/cli/js/K2JsIrCompiler.kt
@@ -61,6 +61,7 @@ import org.jetbrains.kotlin.utils.fileUtils.withReplacedExtensionOrNull
 import org.jetbrains.kotlin.utils.join
 import java.io.File
 import java.io.IOException
+import java.lang.IllegalArgumentException
 
 enum class ProduceKind {
     DEFAULT,  // Determine what to produce based on js-v1 options
@@ -320,7 +321,7 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                 val outputWatFile = outputFile.withReplacedExtensionOrNull(outputFile.extension, "wat")!!
                 outputWatFile.writeText(res.wat)
 
-                val runner = """
+                val esmRunner = """
                     export default WebAssembly.instantiateStreaming(fetch('${outputWasmFile.name}'), { runtime, js_code }).then((it) => {
                         wasmInstance = it.instance;
                         wasmInstance.exports.__init?.();
@@ -330,6 +331,27 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                     });
                 """.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
+                    else -> throw IllegalArgumentException("Unrecognized flavor for the wasm launcher")
+                }
+
                 outputFile.writeText(res.js + "\n" + runner)
                 return OK
             }
diff --git a/compiler/testData/cli/js/jsExtraHelp.out b/compiler/testData/cli/js/jsExtraHelp.out
index f22fff300a3..3a4af35e9d7 100644
--- a/compiler/testData/cli/js/jsExtraHelp.out
+++ b/compiler/testData/cli/js/jsExtraHelp.out
@@ -37,6 +37,7 @@ where advanced options include:
   -Xtyped-arrays             Translate primitive arrays to JS typed arrays
   -Xwasm                     Use experimental WebAssembly compiler backend
   -Xwasm-debug-info          Add debug info to WebAssembly compiled module
+  -Xwasm-launcher=esm|nodejs 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