js exports

Author: MikeInnes

DOCS.md

@@ -71,13 +71,27 @@ $ ./hello.js
 Cacaw, World!
 ```
 
-To run with Node.js directly, you need the `--experimental-wasm-jspi` option.
+You can embed the wasm into a single JS script with `--embed`.
+
+You can run with node directly (prior to v25, you need `--experimental-wasm-jspi`).
 
 ```bash
-$ node --experimental-wasm-jspi hello.js
+$ node hello.js
 Cacaw, World!
 ```
 
+Exported Raven functions are emitted as async JS ones.
+
+```rust
+export { add }
+fn add(a: JSObject, b: JSObject) { Int32(a) + Int32(b) }
+```
+
+```js
+import { add } from './math.js'
+console.log(await add(2, 3))
+```
+
 Profile the compiler:
 
 ```bash

package.json

@@ -18,7 +18,7 @@
   },
   "scripts": {
     "prepare": "npm run build",
-    "build:cli": "esbuild --platform=node --sourcemap --bundle --format=esm --external:commander --outdir=dist/cli src/cli/index.ts src/cli/worker.ts src/cli/exec.ts",
+    "build:cli": "esbuild --platform=node --sourcemap --bundle --format=esm --external:commander --outdir=dist/cli src/cli/index.ts src/cli/worker.ts src/cli/exec.ts src/cli/lib.ts",
     "watch": "npm run build:cli -- --watch",
     "start": "npm run build:cli && node --enable-source-maps --experimental-wasm-jspi dist/cli/index.js",
     "test": "tsc && npm run build:cli && node --max-old-space-size=8192 --enable-source-maps --experimental-wasm-jspi ./node_modules/uvu/bin.js -r tsx test -i '^(?!.*\\.ts$).*'",

src/backend/wasm.ts

@@ -276,13 +276,15 @@ class BatchEmitter implements Emitter {
   seen: Set<string>
   funcs: wasm.Func[]
   imports: wasm.Import[]
+  exports: wasm.Export[]
   constructor(tables: Tables) {
     this.tables = tables
     this.main = []
     this.destructors = []
     this.seen = new Set()
     this.funcs = []
     this.imports = []
+    this.exports = []
   }
 
   clone(): BatchEmitter {
@@ -292,6 +294,7 @@ class BatchEmitter implements Emitter {
     em.seen = new Set(this.seen)
     em.funcs = [...this.funcs]
     em.imports = [...this.imports]
+    em.exports = [...this.exports]
     return em
   }
 
@@ -321,6 +324,11 @@ class BatchEmitter implements Emitter {
     for (const f of this.tables.funcs) this.emitName(calls, f)
     this.destructors.push(func.name)
   }
+
+  export(name: string, as = name) {
+    const existing = this.exports.some(ex => ex.name === name && ex.as === as)
+    if (!existing) this.exports.push(wasm.Export(name, as))
+  }
 }
 
 const refTable = 'jsrefs'
@@ -368,7 +376,8 @@ function wasmmodule(em: BatchEmitter): wasm.Module {
       wasm.Export('cm32p2_memory'),
       wasm.Export(refTable),
       wasm.Export('allocs'),
-      wasm.Export('frees')
+      wasm.Export('frees'),
+      ...em.exports
     ],
     globals: [...em.tables.globals, ...refGlobals].map(g => wasm.Global(...g)),
     tables: moduleTables(em.tables),

src/cli/compile.ts

@@ -1,10 +1,19 @@
 import * as wasm from '../backend/wasm.js'
+import * as types from '../frontend/types.js'
+import * as ast from '../frontend/ast.js'
+import { Binding, MIR, Method } from '../frontend/modules.js'
+import { lowerpattern } from '../frontend/patterns.js'
+import { xcall, xlist, xpart } from '../frontend/lower.js'
 import { Options, withOptions } from '../utils/options.js'
+import { unreachable } from '../utils/ir.js'
+import { reset } from '../utils/cache.js'
 import * as path from 'path'
 import { chmod, mkdir, readFile, writeFile } from 'fs/promises'
 import { spawn, SpawnOptions } from 'node:child_process'
 import { dirname } from './dirname.js'
 import { Compiler } from '../backend/compiler.js'
+import type { Sig } from '../middle/abstract.js'
+import { Def } from '../dwarf/index.js'
 
 export { Compiler, compile, compileJS, exec, load }
 
@@ -20,6 +29,7 @@ interface CompileConfig {
   compiler?: Compiler
   options?: Partial<Options>
   output?: string
+  embed?: boolean
   strip?: boolean
 }
 
@@ -37,21 +47,110 @@ async function compile(file: string, config: CompileConfig = {}): Promise<[Compi
   return [compiler!, wasmPath]
 }
 
+function isJSIdentifier(name: string): boolean {
+  return /^[$A-Z_][0-9A-Z_$]*$/i.test(name)
+}
+
+function buildPaths(file: string, dir: string, output?: string): { js: string, wasm: string } {
+  if (!output) {
+    const base = path.basename(file, path.extname(file))
+    return {
+      js: path.join(dir, `${base}.js`),
+      wasm: path.join(dir, `${base}.wasm`)
+    }
+  }
+  const { dir: outDir, name, ext } = path.parse(output)
+  const wasmBase = ext ? path.join(outDir, name) : output
+  return { js: output, wasm: `${wasmBase}.wasm` }
+}
+
+function exportedFunctions(compiler: Compiler): [string, types.Tag][] {
+  const mod = compiler.pipe.sources.module(types.tag(''))
+  const out: [string, types.Tag][] = []
+  for (const name of [...mod.exports].sort()) {
+    const value = compiler.pipe.defs.resolve_static(new Binding(types.tag(''), name))
+    if (!(value instanceof types.Tag)) continue
+    out.push([name, value])
+  }
+  return out
+}
+
+// TODO better to have a generic means for converting to JS functions. Exported
+// globals can implicitly convert to JS, and we don't need to wrap.
+function libWrapperIR(name: string, f: types.Tag): MIR {
+  const rcall = (code: MIR, fn: types.Tag | Method, args: any[]) => {
+    const arglist = code.push(code.stmt(xlist(...args)))
+    const result = code.push(code.stmt(xcall(fn, arglist)))
+    return code.push(code.stmt(xpart(result, types.Type(1n))))
+  }
+  const code = MIR(Def(name))
+  const args = code.argument(unreachable)
+  const jsargs = rcall(code, types.tag('common.JSObject'), [args])
+  const rvargs = rcall(code, types.tag('common.collect'), [jsargs])
+  const result = code.push(code.stmt(xcall(f, rvargs)))
+  const value = code.push(code.stmt(xpart(result, types.Type(1n))))
+  const jsresult = rcall(code, types.tag('common.js'), [value])
+  code.return(code.push(code.stmt(xpart(jsresult, types.Type(1n)))))
+  return code
+}
+
+function emitSig(compiler: Compiler, em: wasm.BatchEmitter, sig: Sig, main = true): string {
+  reset(compiler.pipe)
+  const func = compiler.pipe.wasm.get(sig)
+  const calls = wasm.calltree(compiler.pipe.wasm, func)
+  const start = em.main.length
+  em.emit(calls, func)
+  if (!main) em.main.length = start
+  return func.name
+}
+
+function jsRuntime(exports: [string, string][], runtime: string, config: { wasmFile?: string, base64?: string, memcheck?: boolean }): string {
+  const { wasmFile, base64, memcheck = false } = config
+  const init = base64
+    ? `\nconst __raven = await __ravenInline(${JSON.stringify(base64)}, ${memcheck})\n`
+    : `\nconst __raven = await __ravenLib(${JSON.stringify(wasmFile)}, ${memcheck})\n`
+  const wrappers = exports.map(([name, wasmName], i) => {
+    const fn = `__raven_fn_${i}`
+    return `const ${fn} = __raven(${JSON.stringify(wasmName)})
+export const ${name} = (...args) => ${fn}(args)`
+  }).join('\n')
+  return `${runtime}${init}${wrappers}\n`
+}
+
 async function compileJS(file: string, config: CompileConfig = {}): Promise<[Compiler, string]> {
-  let { dir = path.dirname(file), compiler, options = {}, output, strip = false } = config
-  const base = path.basename(file, path.extname(file))
-  const jsPath = output ?? path.join(dir, `${base}.js`)
-  await mkdir(path.dirname(jsPath), { recursive: true })
-  await withOptions(options, async () => {
+  let { dir = path.dirname(file), compiler, options = {}, output, embed: inlineWasm = false, strip = false } = config
+  const memcheck = options.memcheck ?? false
+  const paths = buildPaths(file, dir, output)
+  await mkdir(path.dirname(paths.js), { recursive: true })
+  if (!inlineWasm) await mkdir(path.dirname(paths.wasm), { recursive: true })
+  await withOptions({ ...options, memcheck }, async () => {
     compiler ??= await Compiler.create(load)
     const em = await compiler.reload(file)
+    const exports: [string, string][] = []
+    const runtime = await readFile(libPath, 'utf8')
+    const mod = compiler.pipe.sources.module(types.tag(''))
+    for (const [i, [name, fn]] of exportedFunctions(compiler).entries()) {
+      if (!isJSIdentifier(name))
+        throw new Error(`Cannot export ${JSON.stringify(name)} as a JS binding`)
+      const tag = types.tag(`__raven.lib.${i}`)
+      const sig = lowerpattern(ast.List(ast.symbol('args')))
+      const method = mod.method(tag, sig, libWrapperIR(tag.path, fn))
+      const fname = emitSig(compiler, em, [method, types.Ref], false)
+      const wname = `raven.lib.${name}`
+      em.export(fname, wname)
+      exports.push([name, wname])
+    }
     const bytes = wasm.emitwasm(em, strip)
-    const base64 = Buffer.from(bytes).toString('base64')
-    const runtime = await readFile(execPath, 'utf8')
-    await writeFile(jsPath, `${runtime}\nbinary = Buffer.from('${base64}', 'base64')\n`)
-    await chmod(jsPath, 0o755)
+    if (inlineWasm) {
+      const base64 = Buffer.from(bytes).toString('base64')
+      await writeFile(paths.js, jsRuntime(exports, runtime, { base64, memcheck }))
+    } else {
+      await writeFile(paths.wasm, Buffer.from(bytes))
+      await writeFile(paths.js, jsRuntime(exports, runtime, { wasmFile: path.basename(paths.wasm), memcheck }))
+    }
+    await chmod(paths.js, 0o755)
   })
-  return [compiler!, jsPath]
+  return [compiler!, paths.js]
 }
 
 async function run(cmd: string, args: readonly string[] = [], options: SpawnOptions = {}) {
@@ -62,6 +161,7 @@ async function run(cmd: string, args: readonly string[] = [], options: SpawnOpti
   })
 }
 
+const libPath = path.join(dirname, '../../dist/cli/lib.js')
 const execPath = path.join(dirname, '../../dist/cli/exec.js')
 
 async function exec(file: string, args: string[] = [], config?: CompileConfig): Promise<void> {

src/cli/exec.ts

@@ -2,23 +2,20 @@
 import * as fs from 'fs/promises'
 import { loadWasm } from '../backend/support.js'
 
-let binary: string | undefined
+const wasm = await fs.readFile(process.argv[2])
+let { _start, jsrefs, allocs, frees } = await loadWasm(wasm)
+_start = (WebAssembly as any).promising(_start)
 
-async function main() {
-  const wasm = binary ? Buffer.from(binary, 'base64') : await fs.readFile(process.argv[2])
-  let { _start, jsrefs, allocs, frees } = await loadWasm(wasm)
-  _start = (WebAssembly as any).promising(_start)
-  try {
-    await (_start as any)()
-  } catch (e) {
-    console.error(e)
-    process.exit(1)
-  }
-  if (allocs.value !== frees.value)
-    console.warn(`Memory management fault: ${allocs.value} allocs != ${frees.value} frees`)
-  for (let i = 0; i < jsrefs.length; i++)
-    if (jsrefs.get(i) !== null)
-      console.warn("Memory management fault: JSObject")
+try {
+  await (_start as any)()
+} catch (e) {
+  console.error(e)
+  process.exit(1)
 }
 
-setImmediate(() => main())
+if (allocs.value !== frees.value)
+  console.warn(`Memory management fault: ${allocs.value} allocs != ${frees.value} frees`)
+
+for (let i = 0; i < jsrefs.length; i++)
+  if (jsrefs.get(i) !== null)
+    console.warn("Memory management fault: JSObject")

src/cli/index.ts

@@ -119,13 +119,20 @@ async function main() {
     .command('build')
     .description('Compile a Raven source file')
     .argument('<source>', 'Source file to compile')
-    .option('--js', 'Emit JS')
+    .option('--js', 'Emit JS wrapper')
+    .option('--embed', 'Embed WASM into emitted JS')
     .option('-o, --output <file>', 'Rename output file')
     .option('--time', 'Print compiler phase timing information')
-    .action(async (source, { output, js, time }) => {
+    .action(async (source, { output, js, embed, time }) => {
       let { inline, memcheck, strip } = program.optsWithGlobals()
       source = path.resolve(process.cwd(), source)
-      let [compiler] = await (js ? compileJS : compile)(source, { options: { inline, memcheck }, output, strip })
+      const build = js ? compileJS : compile
+      let [compiler] = await build(source, {
+        options: { inline, memcheck: js ? false : memcheck },
+        output,
+        embed: js && !!embed,
+        strip
+      })
       if (time) printTiming(compiler)
     })
 

src/cli/lib.ts

@@ -0,0 +1,34 @@
+#!/usr/bin/env -S node --enable-source-maps --experimental-wasm-jspi
+import * as fs from 'node:fs/promises'
+import { loadWasm } from '../backend/support.js'
+
+export { __ravenInline, __ravenLib }
+
+async function __ravenInline(base64: string, check = false) {
+  return await __ravenLoad(Buffer.from(base64, 'base64'), check)
+}
+
+async function __ravenLib(file: string, check = false) {
+  const wasm = await fs.readFile(new URL(`./${file}`, import.meta.url))
+  return await __ravenLoad(wasm, check)
+}
+
+function checkMemory({ allocs, frees, jsrefs }: any) {
+  if (allocs.value !== frees.value)
+    console.warn(`Memory management fault: ${allocs.value} allocs != ${frees.value} frees`)
+  for (let i = 0; i < jsrefs.length; i++)
+    if (jsrefs.get(i) !== null)
+      console.warn("Memory management fault: JSObject")
+}
+
+async function __ravenLoad(wasm: Uint8Array, check: boolean) {
+  const exports = await loadWasm(wasm)
+  const _start = (WebAssembly as any).promising(exports._start)
+  await _start()
+  if (check) checkMemory(exports)
+  return (name: string) => {
+    const fn = (exports as any)[name]
+    if (typeof fn !== 'function') throw new Error(`Missing Raven export: ${name}`)
+    return (WebAssembly as any).promising(fn)
+  }
+}