[Wasm] Add uninstantiated MJS wrapper

It allows

* Custom imports
* Ability to skip initializer

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

@@ -368,22 +368,22 @@ class K2JsIrCompiler : CLICompiler<K2JSCompilerArguments>() {
                     eliminateDeadDeclarations(allModules, backendContext)
                 }
 
-                val sourceMapFileName = if (configuration.getBoolean(JSConfigurationKeys.SOURCE_MAP)) "$outputName.map" else null
+                val generateSourceMaps = configuration.getBoolean(JSConfigurationKeys.SOURCE_MAP)
 
                 val res = compileWasm(
                     allModules = allModules,
                     backendContext = backendContext,
+                    baseFileName = outputName,
                     emitNameSection = arguments.wasmDebug,
                     allowIncompleteImplementations = arguments.irDce,
                     generateWat = true,
-                    sourceMapFileName = sourceMapFileName
+                    generateSourceMaps = generateSourceMaps
                 )
 
                 writeCompilationResult(
                     result = res,
                     dir = outputDir,
                     fileNameBase = outputName,
-                    sourceMapFileName = sourceMapFileName
                 )
 
                 return OK

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

@@ -33,7 +33,8 @@ import java.io.File
 
 class WasmCompilerResult(
     val wat: String?,
-    val js: String,
+    val jsUninstantiatedWrapper: String,
+    val jsWrapper: String,
     val wasm: ByteArray,
     val sourceMap: String?
 )
@@ -86,10 +87,11 @@ fun compileToLoweredIr(
 fun compileWasm(
     allModules: List<IrModuleFragment>,
     backendContext: WasmBackendContext,
+    baseFileName: String,
     emitNameSection: Boolean = false,
     allowIncompleteImplementations: Boolean = false,
     generateWat: Boolean = false,
-    sourceMapFileName: String? = null,
+    generateSourceMaps: Boolean = false,
 ): WasmCompilerResult {
     val compiledWasmModule = WasmCompiledModuleFragment(backendContext.irBuiltIns)
     val codeGenerator = WasmModuleFragmentGenerator(backendContext, compiledWasmModule, allowIncompleteImplementations = allowIncompleteImplementations)
@@ -105,12 +107,14 @@ fun compileWasm(
         null
     }
 
-    val js = compiledWasmModule.generateJs()
+    val jsUninstantiatedWrapper = compiledWasmModule.generateAsyncJsWrapper("./$baseFileName.wasm")
+    val jsWrapper = generateEsmExportsWrapper("./$baseFileName.uninstantiated.mjs")
 
     val os = ByteArrayOutputStream()
 
+    val sourceMapFileName = "$baseFileName.map".takeIf { generateSourceMaps }
     val sourceLocationMappings =
-        if (sourceMapFileName != null) mutableListOf<SourceLocationMapping>() else null
+        if (generateSourceMaps) mutableListOf<SourceLocationMapping>() else null
 
     val wasmIrToBinary =
         WasmIrToBinary(
@@ -128,7 +132,8 @@ fun compileWasm(
 
     return WasmCompilerResult(
         wat = wat,
-        js = js,
+        jsUninstantiatedWrapper = jsUninstantiatedWrapper,
+        jsWrapper = jsWrapper,
         wasm = byteArray,
         sourceMap = generateSourceMap(backendContext.configuration, sourceLocationMappings)
     )
@@ -167,46 +172,43 @@ private fun generateSourceMap(
     return sourceMapBuilder.build()
 }
 
-fun WasmCompiledModuleFragment.generateJs(): String {
+fun WasmCompiledModuleFragment.generateAsyncJsWrapper(wasmFilePath: String): String {
-    //language=js
+    val jsCodeBody = jsFuns.joinToString(",\n") {
-    val runtime = """
+        "${it.importName.toJsStringLiteral()} : ${it.jsCode}"
-    const externrefBoxes = new WeakMap();
+    }
-    // ref must be non-null
+
-    function tryGetOrSetExternrefBox(ref, ifNotCached) {
+    val jsCodeBodyIndented = jsCodeBody.prependIndent("    ")
-        if (typeof ref !== 'object') return ifNotCached;
-        const cachedBox = externrefBoxes.get(ref);
-        if (cachedBox !== void 0) return cachedBox;
-        externrefBoxes.set(ref, ifNotCached);
-        return ifNotCached;
-    }    """.trimIndent()
 
     val imports = jsModuleImports
         .toList()
         .sorted()
-        .joinToString("\n") {
+        .joinToString("") {
             val moduleSpecifier = it.toJsStringLiteral()
-            "  $moduleSpecifier: await import($moduleSpecifier),"
+            "        $moduleSpecifier: imports[$moduleSpecifier] ?? await import($moduleSpecifier),\n"
         }
 
-    val jsCodeBody = jsFuns.joinToString(",\n") {
+    //language=js
-        "${it.importName.toJsStringLiteral()} : ${it.jsCode}"
+    return """
-    }
+const externrefBoxes = new WeakMap();
-    val jsCodeBodyIndented = jsCodeBody.prependIndent("    ")
+// ref must be non-null
-    val importObject = """
+function tryGetOrSetExternrefBox(ref, ifNotCached) {
-const _import_object = { 
+    if (typeof ref !== 'object') return ifNotCached;
-${imports} 
+    const cachedBox = externrefBoxes.get(ref);
-  js_code: { 
+    if (cachedBox !== void 0) return cachedBox;
-${jsCodeBodyIndented}
+    externrefBoxes.set(ref, ifNotCached);
-  }
+    return ifNotCached;
-};
-"""
-
-    return runtime + importObject
 }
 
-fun generateJsWasmLoader(wasmFilePath: String, externalJs: String): String =
+const js_code = {
-    externalJs + """
+$jsCodeBodyIndented
-    
+}
+
+// Placed here to give access to it from externals (js_code)
+let wasmInstance;
+let require; 
+let wasmExports;
+
+export async function instantiate(imports={}, runInitializer=true) {
     const isNodeJs = (typeof process !== 'undefined') && (process.release.name === 'node');
     const isD8 = !isNodeJs && (typeof d8 !== 'undefined');
     const isBrowser = !isNodeJs && !isD8 && (typeof window !== 'undefined');
@@ -215,8 +217,12 @@ fun generateJsWasmLoader(wasmFilePath: String, externalJs: String): String =
       throw "Supported JS engine not detected";
     }
     
-    let wasmInstance;
+    const wasmFilePath = ${wasmFilePath.toJsStringLiteral()};
-    let require; // Placed here to give access to it from externals (js_code)
+    const importObject = {
+        js_code,
+$imports
+    };
+    
     if (isNodeJs) {
       const module = await import(/* webpackIgnore: true */'node:module');
       require = module.default.createRequire(import.meta.url);
@@ -225,31 +231,40 @@ fun generateJsWasmLoader(wasmFilePath: String, externalJs: String): String =
       const url = require('url');
       const filepath = url.fileURLToPath(import.meta.url);
       const dirpath = path.dirname(filepath);
-      const wasmBuffer = fs.readFileSync(path.resolve(dirpath, '$wasmFilePath'));
+      const wasmBuffer = fs.readFileSync(path.resolve(dirpath, wasmFilePath));
       const wasmModule = new WebAssembly.Module(wasmBuffer);
-      wasmInstance = new WebAssembly.Instance(wasmModule, _import_object);
+      wasmInstance = new WebAssembly.Instance(wasmModule, importObject);
     }
     
     if (isD8) {
-      const wasmBuffer = read('$wasmFilePath', 'binary');
+      const wasmBuffer = read(wasmFilePath, 'binary');
       const wasmModule = new WebAssembly.Module(wasmBuffer);
-      wasmInstance = new WebAssembly.Instance(wasmModule, _import_object);
+      wasmInstance = new WebAssembly.Instance(wasmModule, importObject);
     }
     
     if (isBrowser) {
-      wasmInstance = (await WebAssembly.instantiateStreaming(fetch('$wasmFilePath'), _import_object)).instance;
+      wasmInstance = (await WebAssembly.instantiateStreaming(fetch(wasmFilePath), importObject)).instance;
     }
     
-    const wasmExports = wasmInstance.exports;
+    wasmExports = wasmInstance.exports;
-    wasmExports.__init();
+    if (runInitializer) {
-    export default wasmExports;
+        wasmExports.__init();
-    """.trimIndent()
+    }
+
+    return { instance: wasmInstance,  exports: wasmExports };
+}
+"""
+}
+
+fun generateEsmExportsWrapper(asyncWrapperFileName: String): String = /*language=js */ """
+import { instantiate } from ${asyncWrapperFileName.toJsStringLiteral()};
+export default (await instantiate()).exports;
+"""
 
 fun writeCompilationResult(
     result: WasmCompilerResult,
     dir: File,
-    fileNameBase: String,
+    fileNameBase: String
-    sourceMapFileName: String?
 ) {
     dir.mkdirs()
     if (result.wat != null) {
@@ -257,10 +272,10 @@ fun writeCompilationResult(
     }
     File(dir, "$fileNameBase.wasm").writeBytes(result.wasm)
 
-    val jsWithLoader = generateJsWasmLoader("./$fileNameBase.wasm", result.js)
+    File(dir, "$fileNameBase.uninstantiated.mjs").writeText(result.jsUninstantiatedWrapper)
-    File(dir, "$fileNameBase.mjs").writeText(jsWithLoader)
+    File(dir, "$fileNameBase.mjs").writeText(result.jsWrapper)
 
-    if (sourceMapFileName != null) {
+    if (result.sourceMap != null) {
-        File(dir, sourceMapFileName).writeText(result.sourceMap!!)
+        File(dir, "$fileNameBase.map").writeText(result.sourceMap)
     }
 }

compiler/testData/codegen/boxWasmJsInterop/imperativeWrapperInitialised.kt

@@ -0,0 +1,51 @@
+// TARGET_BACKEND: WASM
+
+// FILE: wasmImport.kt
+import kotlin.wasm.WasmImport
+
+@WasmImport("foo")
+external fun inc1(x: Int): Int
+
+@WasmImport("~!@#\$%^&*()_+`-={}|[]\\\\:\\\";'<>?,./", "inc2")
+external fun inc2(x: Int): Int
+
+@WasmImport("./bar.mjs", "inc3")
+external fun inc3(x: Int): Int
+
+@JsExport
+fun myBox(): String {
+    if (inc1(5) != 6) return "KFail1"
+    if (inc2(5) != 6) return "KFail2"
+    if (inc3(5) != 6) return "KFail3"
+    return "OK"
+}
+
+var initialized: Int = 0
+
+@JsExport
+fun getInitialized(): Int = initialized
+
+fun main() {
+    initialized = 100
+}
+
+// FILE: entry.mjs
+import { instantiate } from "./index.uninstantiated.mjs";
+
+let inc = x => x + 1;
+
+let imports = {
+    "foo": { inc1 : inc },
+    "~!@#\$%^&*()_+\`-={}|[]\\\\:\\\";'<>?,./" : { inc2 : inc },
+    "./bar.mjs" : { inc3 : inc },
+}
+
+let { exports } = await instantiate(imports);
+
+if (exports.getInitialized() !== 100) {
+    throw "Fail1"
+}
+
+if (exports.myBox() != "OK") {
+    throw "Fail2"
+}

compiler/testData/codegen/boxWasmJsInterop/imperativeWrapperUninitialised.kt

@@ -0,0 +1,44 @@
+// TARGET_BACKEND: WASM
+
+// FILE: wasmImport.kt
+import kotlin.wasm.WasmImport
+
+@WasmImport("foo")
+external fun inc(x: Int): Int
+
+@JsExport
+fun myBox(): String {
+    if (inc(5) != 6) return "KFail1"
+    return "OK"
+}
+
+var initialized: Int = 0
+
+@JsExport
+fun getInitialized(): Int = initialized
+
+fun main() {
+    initialized = 100
+}
+
+// FILE: entry.mjs
+import { instantiate } from "./index.uninstantiated.mjs";
+
+let inc = x => x + 1;
+
+let imports = {
+    "foo": { inc },
+}
+
+let { exports } = await instantiate(imports, /*runInitializer=*/false);
+if (exports.getInitialized() !== 0) {
+    throw "Fail1"
+}
+exports.__init();
+if (exports.getInitialized() !== 100) {
+    throw "Fail2"
+}
+
+if (exports.myBox() != "OK") {
+    throw "Fail3"
+}

js/js.tests/test/org/jetbrains/kotlin/js/testOld/BasicWasmBoxTest.kt

@@ -69,6 +69,8 @@ abstract class BasicWasmBoxTest(
             val jsFilesAfter = mutableListOf<String>()
             val mjsFiles = mutableListOf<String>()
 
+            var entryMjs: String? = "test.mjs"
+
             inputFiles.forEach {
                 val name = it.fileName
                 when {
@@ -81,8 +83,13 @@ abstract class BasicWasmBoxTest(
                     name.endsWith(".js") ->
                         jsFilesBefore += name
 
-                    name.endsWith(".mjs") ->
+                    name.endsWith(".mjs") -> {
                         mjsFiles += name
+                        val fileName = File(name).name
+                        if (fileName == "entry.mjs") {
+                            entryMjs = fileName
+                        }
+                    }
                 }
             }
 
@@ -136,10 +143,12 @@ abstract class BasicWasmBoxTest(
             )
 
             val generateWat = debugMode >= DebugMode.DEBUG
+            val baseFileName = "index"
 
             val compilerResult = compileWasm(
                 allModules = allModules,
                 backendContext = backendContext,
+                baseFileName = baseFileName,
                 emitNameSection = true,
                 allowIncompleteImplementations = false,
                 generateWat = generateWat,
@@ -150,6 +159,7 @@ abstract class BasicWasmBoxTest(
             val compilerResultWithDCE = compileWasm(
                 allModules = allModules,
                 backendContext = backendContext,
+                baseFileName = baseFileName,
                 emitNameSection = true,
                 allowIncompleteImplementations = true,
                 generateWat = generateWat,
@@ -189,6 +199,7 @@ abstract class BasicWasmBoxTest(
                     val path = dir.absolutePath
                     println(" ------ $name Wat  file://$path/index.wat")
                     println(" ------ $name Wasm file://$path/index.wasm")
+                    println(" ------ $name JS   file://$path/index.uninstantiated.mjs")
                     println(" ------ $name JS   file://$path/index.mjs")
                     println(" ------ $name Test file://$path/test.mjs")
                     val projectName = "kotlin"
@@ -221,7 +232,7 @@ abstract class BasicWasmBoxTest(
                     )
                 }
 
-                writeCompilationResult(res, dir, "index", sourceMapFileName = null)
+                writeCompilationResult(res, dir, baseFileName)
                 File(dir, "test.mjs").writeText(testJs)
 
                 for (mjsPath: String in mjsFiles) {
@@ -234,7 +245,7 @@ abstract class BasicWasmBoxTest(
                         "--experimental-wasm-gc",
                         *jsFilesBefore.map { File(it).absolutePath }.toTypedArray(),
                         "--module",
-                        "./test.mjs",
+                        "./${entryMjs}",
                         *jsFilesAfter.map { File(it).absolutePath }.toTypedArray(),
                         workingDirectory = dir
                     )

js/js.tests/tests-gen/org/jetbrains/kotlin/js/testOld/wasm/semantics/IrCodegenWasmJsInteropWasmTestGenerated.java

@@ -55,6 +55,16 @@ public class IrCodegenWasmJsInteropWasmTestGenerated extends AbstractIrCodegenWa
         runTest("compiler/testData/codegen/boxWasmJsInterop/functionTypes.kt");
     }
 
+    @TestMetadata("imperativeWrapperInitialised.kt")
+    public void testImperativeWrapperInitialised() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/imperativeWrapperInitialised.kt");
+    }
+
+    @TestMetadata("imperativeWrapperUninitialised.kt")
+    public void testImperativeWrapperUninitialised() throws Exception {
+        runTest("compiler/testData/codegen/boxWasmJsInterop/imperativeWrapperUninitialised.kt");
+    }
+
     @TestMetadata("jsExport.kt")
     public void testJsExport() throws Exception {
         runTest("compiler/testData/codegen/boxWasmJsInterop/jsExport.kt");