utils: add JSFFI utility code

This commit adds JavaScript util code to utils to support the wasm
backend's JSFFI functionality:

- jsffi/post-link.mjs, a post-linker to process the linked wasm module
  and emit a small complement JavaScript ESM module to be used with it
  at runtime
- jsffi/prelude.js, a tiny bit of prelude code as the JavaScript side
  of runtime logic
- jsffi/test-runner.mjs, run the jsffi test cases

Co-authored-by: amesgen <amesgen@amesgen.de>

CODEOWNERS

@@ -65,6 +65,7 @@
 /utils/iserv-proxy/               @angerman @simonmar
 /utils/iserv/                     @angerman @simonmar
 /utils/fs/                        @Phyx
+/utils/jsffi                      @TerrorJack
 
 [WinIO related code]
 /libraries/base/GHC/Event/Windows/                   @Phyx

utils/jsffi/post-link.mjs

+#!/usr/bin/env -S node
+
+// This is the post-linker program that processes a wasm module with
+// ghc_wasm_jsffi custom section and outputs an ESM module that
+// exports a function to generate the ghc_wasm_jsffi wasm imports. It
+// has a simple CLI interface: "./post-link.mjs -i foo.wasm -o
+// foo.js", as well as an exported postLink function that takes a
+// WebAssembly.Module object and returns the ESM module content.
+
+import fs from "node:fs/promises";
+import path from "node:path";
+import util from "node:util";
+
+// Each record in the ghc_wasm_jsffi custom section are 3
+// NUL-terminated strings: name, binder, body. We try to parse the
+// body as an expression and fallback to statements, and return the
+// completely parsed arrow function source.
+function parseRecord([name, binder, body]) {
+  for (const src of [`${binder} => (${body})`, `${binder} => {${body}}`]) {
+    try {
+      new Function(`return ${src};`);
+      return src;
+    } catch (_) {}
+  }
+  throw new Error(`parseRecord ${name} ${binder} ${body}`);
+}
+
+// Parse ghc_wasm_jsffi custom sections in a WebAssembly.Module object
+// and return an array of records.
+function parseSections(mod) {
+  const recs = [];
+  const dec = new TextDecoder("utf-8", { fatal: true });
+  const importNames = new Set(
+    WebAssembly.Module.imports(mod)
+      .filter((i) => i.module === "ghc_wasm_jsffi")
+      .map((i) => i.name),
+  );
+  for (const buf of WebAssembly.Module.customSections(mod, "ghc_wasm_jsffi")) {
+    const ba = new Uint8Array(buf);
+    let strs = [];
+    for (let l = 0, r; l < ba.length; l = r + 1) {
+      r = ba.indexOf(0, l);
+      strs.push(dec.decode(ba.subarray(l, r)));
+      if (strs.length === 3) {
+        if (importNames.has(strs[0])) {
+          recs.push(strs);
+        }
+        strs = [];
+      }
+    }
+  }
+  return recs;
+}
+
+
+export async function postLink(mod) {
+  let src = await fs.readFile(path.join(import.meta.dirname, "prelude.js"), {
+    encoding: "utf-8",
+  });
+  src = `${src}\nexport default (__exports) => {`;
+  src = `${src}\nconst __ghc_wasm_jsffi_jsval_manager = new JSValManager();`;
+  src = `${src}\nconst __ghc_wasm_jsffi_finalization_registry = new FinalizationRegistry(sp => __exports.rts_freeStablePtr(sp));`;
+  src = `${src}\nreturn {`;
+  src = `${src}\nnewJSVal: (v) => __ghc_wasm_jsffi_jsval_manager.newJSVal(v),`;
+  src = `${src}\ngetJSVal: (k) => __ghc_wasm_jsffi_jsval_manager.getJSVal(k),`;
+  src = `${src}\nfreeJSVal: (k) => __ghc_wasm_jsffi_jsval_manager.freeJSVal(k),`;
+  src = `${src}\nscheduleWork: () => setImmediate(__exports.rts_schedulerLoop),`;
+  for (const rec of parseSections(mod)) {
+    src = `${src}\n${rec[0]}: ${parseRecord(rec)},`;
+  }
+  return `${src}\n};\n};\n`;
+}
+
+function isMain() {
+  return import.meta.filename === process.argv[1];
+}
+
+async function main() {
+  const { input, output } = util.parseArgs({
+    options: {
+      input: {
+        type: "string",
+        short: "i",
+      },
+      output: {
+        type: "string",
+        short: "o",
+      },
+    },
+  }).values;
+
+  await fs.writeFile(
+    output,
+    await postLink(await WebAssembly.compile(await fs.readFile(input)))
+  );
+}
+
+if (isMain()) {
+  await main();
+}

utils/jsffi/prelude.js

+// This file implements the JavaScript runtime logic for Haskell
+// modules that use JSFFI. It is not an ESM module, but the template
+// of one; the post-linker script will copy all contents into a new
+// ESM module.
+
+// Manage a mapping from unique 32-bit ids to actual JavaScript
+// values.
+class JSValManager {
+  #lastk = 0;
+  #kv = new Map();
+
+  constructor() {}
+
+  // Maybe just bump this.#lastk? For 64-bit ids that's sufficient,
+  // but better safe than sorry in the 32-bit case.
+  #allocKey() {
+    let k = this.#lastk;
+    while (true) {
+      if (!this.#kv.has(k)) {
+        this.#lastk = k;
+        return k;
+      }
+      k = (k + 1) | 0;
+    }
+  }
+
+  newJSVal(v) {
+    const k = this.#allocKey();
+    this.#kv.set(k, v);
+    return k;
+  }
+
+  // A separate has() call to ensure we can store undefined as a value
+  // too. Also, unconditionally check this since the check is cheap
+  // anyway, if the check fails then there's a use-after-free to be
+  // fixed.
+  getJSVal(k) {
+    if (!this.#kv.has(k)) {
+      throw new WebAssembly.RuntimeError(`getJSVal(${k})`);
+    }
+    return this.#kv.get(k);
+  }
+
+  // Check for double free as well.
+  freeJSVal(k) {
+    if (!this.#kv.delete(k)) {
+      throw new WebAssembly.RuntimeError(`freeJSVal(${k})`);
+    }
+  }
+}
+
+// A simple & fast setImmediate() implementation for browsers. It's
+// not a drop-in replacement for node.js setImmediate() because:
+// 1. There's no clearImmediate(), and setImmediate() doesn't return
+//    anything
+// 2. There's no guarantee that callbacks scheduled by setImmediate()
+//    are executed in the same order (in fact it's the opposite lol),
+//    but you are never supposed to rely on this assumption anyway
+class SetImmediate {
+  #fs = [];
+  #mc = new MessageChannel();
+
+  constructor() {
+    this.#mc.port1.addEventListener("message", () => {
+      this.#fs.pop()();
+    });
+    this.#mc.port1.start();
+  }
+
+  setImmediate(cb, ...args) {
+    this.#fs.push(() => cb(...args));
+    this.#mc.port2.postMessage(undefined);
+  }
+}
+
+// The actual setImmediate() to be used. This is a ESM module top
+// level binding and doesn't pollute the globalThis namespace.
+let setImmediate;
+if (globalThis.setImmediate) {
+  // node.js, bun
+  setImmediate = globalThis.setImmediate;
+} else {
+  try {
+    // deno
+    setImmediate = (await import("node:timers")).setImmediate;
+  } catch {
+    // browsers
+    const sm = new SetImmediate();
+    setImmediate = (cb, ...args) => sm.setImmediate(cb, ...args);
+  }
+}

utils/jsffi/test-runner.mjs

+#!/usr/bin/env -S node --expose-gc
+
+import fs from "node:fs/promises";
+import path from "node:path";
+import { WASI } from "wasi";
+import { postLink } from "./post-link.mjs";
+
+// The ESM code returned by post-linker doesn't need to be written to
+// a temporary file first. It can be directly imported from a
+// base64-encoded data URL.
+
+function jsToDataURL(src) {
+  return `data:text/javascript;base64,${Buffer.from(src).toString("base64")}`;
+}
+
+async function evalJSModuleDefault(src) {
+  return (await import(jsToDataURL(src))).default;
+}
+
+const wasm_path = path.resolve(process.argv[2]);
+const js_path = path.join(
+  path.dirname(wasm_path),
+  `${path.basename(wasm_path, ".wasm")}.mjs`
+);
+
+const wasm_module = await WebAssembly.compile(await fs.readFile(wasm_path));
+const js_stub_src = await postLink(wasm_module);
+const wasm_import_factory = await evalJSModuleDefault(js_stub_src);
+
+// Yes, you can pass +RTS and other command line flags to Haskell via
+// "./test-runner.mjs foo.wasm +RTS ..."
+const wasi = new WASI({
+  version: "preview1",
+  args: process.argv.slice(2),
+  env: { PATH: "", PWD: process.cwd() },
+  preopens: { "/": "/" },
+});
+
+// Poor man's tying-the-knot
+let __exports = {};
+
+const wasm_instance = await WebAssembly.instantiate(wasm_module, {
+  ghc_wasm_jsffi: wasm_import_factory(__exports),
+  wasi_snapshot_preview1: wasi.wasiImport,
+});
+
+// Do this immediately before you _initialize()
+Object.assign(__exports, wasm_instance.exports);
+
+// This calls _initialize(). Other wasi implementations may differ,
+// always check their doc/src to be sure
+wasi.initialize(wasm_instance);
+
+const k = (await import(js_path)).default;
+await k(__exports);