From 69c97688c0f131a60b2d90bcd7a6b17223be84b3 Mon Sep 17 00:00:00 2001
From: Igor Yakovlev <igor.yakovlev@jetbrains.com>
Date: Wed, 2 Aug 2023 15:05:49 +0200
Subject: [PATCH] [Wasi] Add default wasi ejs runner (for nodejs)

---
 .../jetbrains/kotlin/backend/wasm/compiler.kt | 48 +++++++++++++++----
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/compiler.kt b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/compiler.kt
index 2a39bd1da34..bcae046f3d6 100644
--- a/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/compiler.kt
+++ b/compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/compiler.kt
@@ -22,6 +22,8 @@ import org.jetbrains.kotlin.ir.declarations.IrFactory
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.util.ExternalDependenciesGenerator
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
+import org.jetbrains.kotlin.js.config.JSConfigurationKeys
+import org.jetbrains.kotlin.js.config.WasmTarget
 import org.jetbrains.kotlin.js.sourceMap.SourceFilePathResolver
 import org.jetbrains.kotlin.js.sourceMap.SourceMap3Builder
 import org.jetbrains.kotlin.name.FqName
@@ -34,7 +36,7 @@ import java.io.File
 
 class WasmCompilerResult(
     val wat: String?,
-    val jsUninstantiatedWrapper: String,
+    val jsUninstantiatedWrapper: String?,
     val jsWrapper: String,
     val wasm: ByteArray,
     val sourceMap: String?
@@ -109,12 +111,6 @@ fun compileWasm(
         null
     }
 
-    val jsUninstantiatedWrapper = compiledWasmModule.generateAsyncJsWrapper(
-        "./$baseFileName.wasm",
-        backendContext.jsModuleAndQualifierReferences
-    )
-    val jsWrapper = generateEsmExportsWrapper("./$baseFileName.uninstantiated.mjs")
-
     val os = ByteArrayOutputStream()
 
     val sourceMapFileName = "$baseFileName.map".takeIf { generateSourceMaps }
@@ -134,6 +130,18 @@ fun compileWasm(
     wasmIrToBinary.appendWasmModule()
 
     val byteArray = os.toByteArray()
+    val jsUninstantiatedWrapper: String?
+    val jsWrapper: String
+    if (backendContext.configuration.get(JSConfigurationKeys.WASM_TARGET, WasmTarget.JS) == WasmTarget.JS) {
+        jsUninstantiatedWrapper = compiledWasmModule.generateAsyncJsWrapper(
+            "./$baseFileName.wasm",
+            backendContext.jsModuleAndQualifierReferences
+        )
+        jsWrapper = generateEsmExportsWrapper("./$baseFileName.uninstantiated.mjs")
+    } else {
+        jsUninstantiatedWrapper = null
+        jsWrapper = compiledWasmModule.generateAsyncWasiWrapper("./$baseFileName.wasm")
+    }
 
     return WasmCompilerResult(
         wat = wat,
@@ -177,6 +185,28 @@ private fun generateSourceMap(
     return sourceMapBuilder.build()
 }
 
+fun WasmCompiledModuleFragment.generateAsyncWasiWrapper(wasmFilePath: String): String = """
+import { WASI } from 'wasi';
+import { argv, env } from 'node:process';
+
+const wasi = new WASI({ version: 'preview1', args: argv, env, });
+
+const module = await import(/* webpackIgnore: true */'node:module');
+const require = module.default.createRequire(import.meta.url);
+const fs = require('fs');
+const path = require('path');
+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 wasmModule = new WebAssembly.Module(wasmBuffer);
+const wasmInstance = new WebAssembly.Instance(wasmModule, wasi.getImportObject());
+
+wasi.initialize(wasmInstance);
+
+export default wasmInstance.exports;
+"""
+
 fun WasmCompiledModuleFragment.generateAsyncJsWrapper(
     wasmFilePath: String,
     jsModuleAndQualifierReferences: Set<JsModuleAndQualifierReference>
@@ -341,7 +371,9 @@ fun writeCompilationResult(
     }
     File(dir, "$fileNameBase.wasm").writeBytes(result.wasm)
 
-    File(dir, "$fileNameBase.uninstantiated.mjs").writeText(result.jsUninstantiatedWrapper)
+    if (result.jsUninstantiatedWrapper != null) {
+        File(dir, "$fileNameBase.uninstantiated.mjs").writeText(result.jsUninstantiatedWrapper)
+    }
     File(dir, "$fileNameBase.mjs").writeText(result.jsWrapper)
 
     if (result.sourceMap != null) {