jz 0.8.0 → 0.9.0

package/index.js

@@ -16,12 +16,13 @@
  *     ↓  compile — drives per-function emit, interleaves analysis (locals/valTypes/captures/
  *        narrowing fixpoint) with IR generation via the emitter table (src/compile/emit.js).
  *        Writes: `ctx.func.valTypes`/`.locals`, `ctx.types.*`, `ctx.runtime.*`, `ctx.core.includes`.
- *        Also calls optimizeFunc (src/optimize/index.js): `hoistPtrType` + fused peephole/inline/memarg walk.
+ *        The emit phase (src/wat/assemble.js optimizeModule) then runs jz's ONLY optimizer pass —
+ *        optimizeFunc (src/optimize/index.js): `hoistPtrType` + fused peephole/inline/memarg walk +
+ *        auto-vectorization. All lowering, incl. SIMD, happens here — BEFORE watr.
  *   WAT IR: watr S-expression `['module', ...sections]`, every instruction node carries `.type`.
- *     ↓  watOptimize (opt-out via opts.optimize=false) — CSE, DCE, const folding at WAT level
- *     ↓  optimizeFunc 2nd pass — re-folds rebox/unbox roundtrips that watOptimize's inliner
- *        re-introduces at inline boundaries (caller's boxPtrIR meets callee's
- *        i32.wrap_i64(i64.reinterpret_f64 __env)). watr's peephole doesn't cover this.
+ *     ↓  watOptimize (opt-out via opts.optimize=false) — the SOLE, FINAL optimizer: CSE, DCE, const
+ *        fold, inline, coalesce. Runs ONCE, as a fixpoint. No jz pass touches WAT after it (bar the
+ *        stable-global-offset hoist, a phase-2 watr-migration candidate). See .work/research.md.
  *     ↓  watrPrint (opts.wat=true) → WAT text, or watrCompile → Uint8Array binary
  *
  * # State
@@ -42,17 +43,17 @@
 
 import { parse } from './src/parse.js'
 import watrCompile from "watr/compile";
+import { snapshotInit } from "./src/snapshot.js";
 import watrPrint from "watr/print";
 import watOptimize from "watr/optimize";
-import { appendLateStdlib } from './src/wat/assemble.js'
 import { ctx, reset, err, initWarnings, assertCtxInvariants } from './src/ctx.js'
 import prepare, { GLOBALS } from './src/prepare/index.js'
 import { liftIIFEs } from './src/prepare/lift-iife.js'
+import { preEval } from './src/prepare/pre-eval.js'
 import compile from './src/compile/index.js'
 import { resetProgramFactsCache } from './src/compile/program-facts.js'
 import { emit, emitter, emitVoid as flat, emitBlockBody as body, emitBoolStr as bool, emitIndex as idx, buildArrayWithSpreads as spread } from './src/compile/emit.js'
-import { optimizeFunc, foldStrDispatchF64, collectVolatileGlobals, collectReachableGlobalWrites, hoistGlobalPtrOffset, stablePtrGlobalNames, resolveOptimize, SIMD_PINNED } from './src/optimize/index.js'
-import { findBodyStart } from './src/ir.js'
+import { collectReachableGlobalWrites, hoistGlobalPtrOffset, stablePtrGlobalNames, resolveOptimize, SIMD_PINNED } from './src/optimize/index.js'
 import { VAL } from './src/reps.js'
 import jzify from './jzify/index.js'
 import { T } from './src/ast.js'
@@ -198,6 +199,80 @@ const appendFunctionNames = (wasm, module) => {
  */
 jz.memory = enhanceMemory
 
+/**
+ * jz.pool(source, opts) — SPMD worker pool over ONE shared memory (Workers v1).
+ *
+ * Compiles `source` with { sharedMemory: true }, instantiates it on the main
+ * thread AND in `threads` node worker_threads, all linked to the same
+ * WebAssembly.Memory({ shared: true }). Kernels coordinate through shared
+ * typed arrays + Atomics (module/atomics.js); strings/objects stay
+ * thread-local by the v1 contract.
+ *
+ * Worker instances link `env.memory` ONLY — a kernel needing other host
+ * imports fails instantiation loudly (pure-compute contract).
+ *
+ *   const p = await jz.pool(src, { threads: 4, pages: 16, maxPages: 256 })
+ *   p.exports.setup(...)                    // main-thread instance
+ *   await p.run('tile')                     // every worker: tile(workerIndex, threads)
+ *   await p.run('tile', a, b)               // extra args broadcast to each worker
+ *   p.memory                                // the shared, jz-enhanced memory
+ *   await p.terminate()
+ *
+ * @param {string} source - jz source (compiled once; module shared with workers)
+ * @param {object} [opts] - { threads=4, pages=16, maxPages=16384, ...compile opts }
+ * @returns {Promise<{exports, memory, module, threads, run, terminate}>}
+ */
+jz.pool = async function pool(source, opts = {}) {
+  const { threads = 4, pages = 16, maxPages = 16384, ...rest } = opts
+  const memory = new WebAssembly.Memory({ initial: pages, maximum: maxPages, shared: true })
+  const main = jz(source, { ...rest, sharedMemory: true, memory })
+  // Computed specifier keeps bundlers (esbuild web dist) from resolving the
+  // node builtin statically; browsers get a clear v1 error instead.
+  const { Worker } = await import('node:' + 'worker_threads').catch(() => {
+    throw new Error('jz.pool v1 runs on node worker_threads; browser Worker support is a follow-up')
+  })
+  // The worker shim honors the jz:i64exp lane map (exact-bits ABI): an i64 lane
+  // takes a BigInt — scalars convert to their f64 bits, boxed pointers (BigInt
+  // args from jz.memory.* on the main thread) pass through; i64 results
+  // reinterpret back to numbers (v1 kernels return scalars).
+  const workerSrc = `
+    const { parentPort, workerData } = require('node:worker_threads')
+    const inst = new WebAssembly.Instance(workerData.module, { env: { memory: workerData.memory } })
+    const dv = new DataView(new ArrayBuffer(8))
+    const bits = (x) => (dv.setFloat64(0, x), dv.getBigUint64(0))
+    const unbits = (b) => (dv.setBigUint64(0, BigInt.asUintN(64, b)), dv.getFloat64(0))
+    const lanes = new Map()
+    for (const s of WebAssembly.Module.customSections(workerData.module, 'jz:i64exp'))
+      try { for (const e of JSON.parse(new TextDecoder().decode(s))) lanes.set(e.name, e) } catch {}
+    parentPort.on('message', (m) => {
+      try {
+        const sig = lanes.get(m.fn), p = new Set(sig?.p || [])
+        const args = m.args.map((x, i) => p.has(i) && typeof x === 'number' ? bits(x) : x)
+        let r = inst.exports[m.fn](...args)
+        if (sig?.r && typeof r === 'bigint') r = unbits(r)
+        parentPort.postMessage({ id: m.id, r })
+      } catch (e) { parentPort.postMessage({ id: m.id, e: String(e) }) }
+    })
+    parentPort.postMessage({ ready: true })`
+  const workers = Array.from({ length: threads }, () =>
+    new Worker(workerSrc, { eval: true, workerData: { module: main.module, memory } }))
+  await Promise.all(workers.map(w => new Promise((res, rej) => {
+    w.once('message', res); w.once('error', rej)
+  })))
+  let seq = 0
+  const call = (w, fn, args) => new Promise((res, rej) => {
+    const id = seq++
+    const on = (m) => { if (m.id !== id) return; w.off('message', on); m.e ? rej(new Error(m.e)) : res(m.r) }
+    w.on('message', on)
+    w.postMessage({ id, fn, args })
+  })
+  return {
+    exports: main.exports, memory: main.memory, module: main.module, threads,
+    run: (fn, ...args) => Promise.all(workers.map((w, i) => call(w, fn, [i, threads, ...args]))),
+    terminate: () => Promise.all(workers.map(w => w.terminate())),
+  }
+}
+
 /**
  * Compile jz source to WASM binary or WAT text. Low-level — no instantiation.
  * @param {string} code - jz source
@@ -212,6 +287,9 @@ jz.memory = enhanceMemory
  * @param {number} [opts.maxMemory] - Maximum memory pages — emits a ceiling on the
  *   memory type so growth traps past it (sandbox cap). Must be ≥ the initial size.
  *   Default: unbounded.
+ * @param {boolean} [opts.sharedMemory] - Import `env.memory` as a SHARED memory (wasm
+ *   threads): atomic heap bump, `shared` memtype (max defaults to the 4 GiB ceiling).
+ *   Link with `new WebAssembly.Memory({ initial, maximum, shared: true })`.
  * @param {boolean} [opts.importMemory] - Import `env.memory` instead of exporting an
  *   owned memory. For embedding into a host that provides the memory.
  * @param {boolean} [opts.alloc=true] - Export raw allocator helpers
@@ -385,8 +463,6 @@ const detectOptimizeConfig = (ast, code) => {
     cfg.scalarTypedLoopUnroll = s.maxLoopDepth > 1 ? 8 : 16
     cfg.scalarTypedNestedUnroll = s.maxLoopDepth > 1 ? 32 : 128
   }
-  // String-heavy: ensure pool sorting is on (already default, but explicit).
-  if (s.stringLiteralCount > 30) cfg.sortStrPoolByFreq = true
   // Closure-heavy: ptr hoists pay off.
   if (s.closureCount > 4) {
     cfg.hoistPtrType = true
@@ -432,6 +508,11 @@ const setupCtx = (code, opts) => {
   if (typeof opts.memory === 'number') ctx.memory.pages = opts.memory
   else if (opts.memory) ctx.memory.shared = true
   if (opts.importMemory) ctx.memory.shared = true   // import env.memory instead of exporting own
+  // True cross-thread sharing (Workers v1): import env.memory declared with the
+  // wasm `shared` memtype and switch the heap bump to atomic RMW. Distinct from
+  // importMemory — a plain imported (non-shared) Memory must NOT declare shared
+  // or linking fails in the other direction.
+  if (opts.sharedMemory) { ctx.memory.shared = true; ctx.memory.atomic = true }
   if (opts.maxMemory != null) {
     if (!Number.isInteger(opts.maxMemory) || opts.maxMemory < 1)
       err(`opts.maxMemory must be a positive integer page count (each page is 64 KiB); got ${opts.maxMemory}`)
@@ -517,8 +598,12 @@ const jzCompileInner = (code, opts = {}) => {
   // path. A no-op when there are none.
   parsed = time('liftIIFE', () => liftIIFEs(parsed))
   if (!opts.strict) parsed = time('jzify', () => jzify(parsed))
-  const ast = time('prepare', () => prepare(parsed))
+  let ast = time('prepare', () => prepare(parsed))
   assertCtxInvariants('post-prepare')
+  // preEval: fold every statically-evaluable construct (numeric/string/bool chains,
+  // pure Math.* calls, zero-arg pure calls incl. lift-iife's IIFEs) down to literals,
+  // over the prepared AST + every ctx.func.list body, before compile ever sees them.
+  ast = time('preEval', () => preEval(ast))
 
   // Auto-detect optimization tuning from source characteristics when the user
   // hasn't provided any optimize option. detectOptimizeConfig has two *live*
@@ -609,35 +694,65 @@ const jzCompileInner = (code, opts = {}) => {
     if (watrOpts === true) watrOpts = { inline: 'simd' }
     else if (typeof watrOpts === 'object' && watrOpts.inline === undefined) watrOpts.inline = 'simd'
   }
+  // Wrapper elision (speed tier): dissolve adapter frames — closure-ABI
+  // trampolines for functions-as-values, thin dispatch heads like
+  // __dyn_get_expr — into their target at the wrapper site. The recorded
+  // "table entries native against ftN" restructure, done mechanically in the
+  // optimizer instead of per-case emit surgery (one wrapper = one duplicated
+  // target; treeshake collects whichever copy ends up unreferenced).
+  if (cfg.inlineFns) {
+    if (typeof watrOpts === 'object' && watrOpts.inlineWrappers === undefined) watrOpts.inlineWrappers = true
+  }
+  // watr LICM (speed tier): hoist loop-invariant pure arithmetic AFTER watr's inlining exposes it
+  // (the raytrace per-iteration NaN/Inf guard pairs). Mechanical — lives in watr, jz just enables it.
+  if (cfg.watrLicm) {
+    if (watrOpts === true) watrOpts = { licm: true }
+    else if (typeof watrOpts === 'object' && watrOpts.licm === undefined) watrOpts.licm = true
+  }
   // guardRefine folds jz's NaN-box tag reads — default-off in watr (general WAT
   // never matches it), so jz always enables it explicitly.
   if (watrOpts === true) watrOpts = { guardRefine: true }
   else if (typeof watrOpts === 'object' && watrOpts.guardRefine === undefined) watrOpts.guardRefine = true
+  // watr's ifset (one-armed conditional update → select) rides jz's select
+  // tiering: an unconditional-update trade belongs to the speed tier exactly
+  // like boolConvertToSelect. jz's DEFAULT tier also passes watrProfile:'speed'
+  // (for the outline/tailmerge shape), which would enable ifset via the
+  // profile — the explicit flag here overrides it in both directions.
+  if (typeof watrOpts === 'object' && watrOpts.ifset === undefined)
+    watrOpts.ifset = cfg.boolConvertToSelect === true || cfg.watrIfset === true
+  // jz's promise is runtime speed, but watr's OWN profile default leans size — outline/
+  // tailmerge/rettail fold repeated sequences into out-of-line calls (measured 1.433→1.316
+  // on the self-host kernel with them off, watr ≥5.2.0). Every speed-tier preset carries
+  // watrProfile:'speed' (src/optimize/index.js LEVEL_PRESETS); the 'size' preset keeps
+  // watr's size-leaning default. An explicit user profile always wins.
+  if (watrOpts.profile === undefined && cfg.watrProfile) watrOpts.profile = cfg.watrProfile
+  // Speed tier waives watr's size-revert guard (two full binary encodes per
+  // optimize — its costliest fixed overhead): inlining growth is the shape this
+  // tier asks for, and jz's own perf/size gates own the size budget. Level 2
+  // (default) and 'size' keep the guard — watr's never-inflate contract stands
+  // where the user didn't opt into speed-for-size.
+  if (watrOpts.guard === undefined && cfg.watrGuard === false) watrOpts.guard = false
   // Pin jz's scalar transcendentals (the PPC_CALL2 keys the auto-vectorizer rewrites to f64x2
   // mirrors) so watr's inline passes don't dissolve the call nodes the lift needs. The protection
   // policy lives here in jz — watr just honours the `pin` list (no jz names hardcoded in watr).
   if (cfg.watr) {
     if (watrOpts === true) watrOpts = {}
     watrOpts.pin = watrOpts.pin ? [...watrOpts.pin, ...SIMD_PINNED] : SIMD_PINNED
+    // Boundary-wrapped functions (`$name`/`$name$exp`) the vectorizer just lifted: keep
+    // watr's default inlineOnce from merging $name into $name$exp, which would otherwise
+    // rename every $__ppc*/$__esc*/... local & label the lift produced with its own
+    // $__inl{N}_ prefix, burying the marker's leading `$` (cosmetic-only cost: one
+    // un-eliminated wrapper call per top-level invocation, not per loop iteration).
+    if (cfg._vectorizedFnNames?.size) watrOpts.pin = [...watrOpts.pin, ...cfg._vectorizedFnNames]
   }
-  // Capture the per-func alias facts (param-distinctness) the emit phase stamped as JS
-  // expandos — watOptimize round-trips the module through watr's printer/parser, which
-  // rebuilds plain arrays and DROPS every non-index property. Without re-attaching, the
-  // 'post' leaf passes (splitLoopPrivateScratch, hoistInvariantLoop) lose the proven
-  // alias model and fall back to weaker/heuristic invariance — the exact reason raytrace's
-  // read-only sphere loads couldn't be SOUNDLY proven loop-invariant after watr ran.
-  const aliasFacts = cfg.watr ? new Map(module.filter(n => Array.isArray(n) && n[0] === 'func' && n.distinctParams).map(n => [n[1], n.distinctParams])) : null
   const optimized = cfg.watr ? time('watOptimize', () => watOptimize(module, watrOpts)) : module
-  if (aliasFacts?.size) for (const n of optimized) if (Array.isArray(n) && n[0] === 'func' && aliasFacts.has(n[1])) n.distinctParams = aliasFacts.get(n[1])
-  // Stable-pointee module globals: resolve the __ptr_offset once per function.
-  // Never-forwarding kinds — every PTR tag outside __ptr_offset's forwarding
-  // set {ARRAY, HASH, SET, MAP} — give the same offset for the same bits, so
-  // the snapshot only needs the global's VALUE stable through the function:
-  // the reachable-writes call graph proves that precisely. Independent of watr
-  // (the auto-config turns watr off for large sources — exactly the
-  // module-global DSP-state programs this pass exists for: rfft, diffusion);
-  // when watr DID run, it goes before the post leaf passes so the snapped
-  // local participates in hoistAddrBase/cseScalarLoad.
+  // Stable-pointee module globals: resolve the __ptr_offset once per function. Never-forwarding
+  // kinds — every PTR tag outside __ptr_offset's forwarding set {ARRAY, HASH, SET, MAP} — give the
+  // same offset for the same bits, so the snapshot only needs the global's VALUE stable through the
+  // function: the reachable-writes call graph proves that precisely. This is the ONE jz pass still
+  // touching WAT after watr; it exists because watr has no stable-global-offset hoist of its own
+  // (the module-global DSP programs it serves — rfft, diffusion — run with watr auto-OFF anyway).
+  // TODO(phase-2): migrate into watr so the "watr is the sole optimizer" invariant holds fully.
   if (cfg.hoistGlobalPtrOffset !== false) {
     const stableGlobals = stablePtrGlobalNames()
     if (stableGlobals.size) {
@@ -646,52 +761,20 @@ const jzCompileInner = (code, opts = {}) => {
       for (const node of funcs) hoistGlobalPtrOffset(node, stableGlobals, reach)
     }
   }
-  // Final peephole pass: watOptimize's inliner can re-introduce rebox/unbox at boundaries
-  // (e.g. inlined closure body's `i32.wrap_i64 (i64.reinterpret_f64 __env)` next to caller's
-  // `boxPtrIR(g)` rebox). Our fusedRewrite folds these, watr's peephole doesn't.
-  // Only valuable to re-run when watr ran (watr is what re-introduces the boundaries).
-  if (cfg.watr) {
-    // Build global name→type map from ctx.scope.globalTypes for promoteGlobals
-    const globalTypesMap = ctx.scope.globalTypes ? new Map([...ctx.scope.globalTypes].map(([k, v]) => [`$${k}`, v])) : null
-    // Build pure-function map for Phase 2 user-function inline in tryPerPixelColor.
-    // A function is "pure for SIMD inline" if its body contains no side effects:
-    // no global.set, no memory stores, no calls outside $math.*.
-    // foldStrDispatchF64 is run first (idempotent) so the purity check sees the
-    // folded body — dead __is_str_key dispatch would otherwise look impure.
-    if (cfg.vectorizeLaneLocal === true) {
-      const allFuncs = optimized.filter(node => Array.isArray(node) && node[0] === 'func')
-      const pureFuncMap = new Map()
-      const hasSideEffect = (node) => {
-        if (!Array.isArray(node)) return false
-        const op = node[0]
-        if (op === 'global.set') return true
-        if (typeof op === 'string' && (op.endsWith('.store') || op.startsWith('memory.'))) return true
-        if (op === 'call' && typeof node[1] === 'string' && !node[1].startsWith('$math.')) return true
-        if (op === 'call_indirect' || op === 'call_ref') return true
-        return node.some((c, i) => i > 0 && hasSideEffect(c))
-      }
-      for (const fn of allFuncs) {
-        const name = fn[1]
-        if (typeof name !== 'string' || name.startsWith('$math.') || name.startsWith('$__')) continue
-        // Fold dead str-dispatch blocks so purity check sees the clean form.
-        foldStrDispatchF64(fn)
-        const bodyStart = findBodyStart(fn)
-        if (bodyStart < 0) continue
-        let pure = true
-        for (let i = bodyStart; i < fn.length; i++) if (hasSideEffect(fn[i])) { pure = false; break }
-        if (pure) pureFuncMap.set(name, fn)
-      }
-      if (pureFuncMap.size) cfg._pureFuncMap = pureFuncMap
-    }
-    time('watrReopt', () => {
-      const funcs = optimized.filter(node => Array.isArray(node) && node[0] === 'func')
-      const volatileGlobals = collectVolatileGlobals(funcs)
-      const reach = collectReachableGlobalWrites(funcs)
-      for (const node of funcs) optimizeFunc(node, cfg, globalTypesMap, volatileGlobals, 'post', reach)
-    })
-    // The 'post' lane vectorizer can inject stdlib calls (e.g. the f64x2 trig mirror $math.cos2)
-    // absent from the already-pulled+treeshaken module — append any now-referenced helper body.
-    appendLateStdlib(optimized)
+  // NO post-watr optimizer. jz does ALL lowering — including auto-vectorization — BEFORE watr
+  // (src/wat/assemble.js optimizeModule → optimizeFunc 'pre'); watr is the sole optimizer and runs
+  // exactly ONCE, as a fixpoint. Re-running jz's leaf passes on watr's output (the former 'post'
+  // phase) both violated that architecture AND miscompiled: its propagate/fold sweep dropped a
+  // reassigned-param `local.tee` write and corrupted the divergent-escape SIMD frame (mandelbrot).
+  // The f64x2 stdlib mirrors the pre-watr vectorizer injects are appended in the emit phase
+  // (optimizeModule's appendLateStdlib); nothing post-watr needs to touch the module.
+  // Pre-eval tier 3 — module-init snapshotting (src/snapshot.js): run __start once
+  // NOW, bake the post-init heap image + global values into the artifact, delete
+  // __start. Opt-in (optimize.snapshotInit); declined cleanly (dynamically-proven
+  // hermeticity) when init touches the host, loops forever, or memory is shared.
+  if (cfg.snapshotInit) {
+    const took = time('snapshotInit', () => snapshotInit(optimized, watrCompile))
+    if (!took && opts.warnings) ctx.warnings?.push?.({ code: 'snapshot-declined', message: 'init snapshot declined (host-touching, timer, or shared-memory init) — compiled without it' })
   }
   try {
     if (opts.wat) {
@@ -805,3 +888,11 @@ export const compileModule = (code, opts = {}) => {
 }
 
 export { instantiateRuntime as instantiate }
+
+/**
+ * jzify as a standalone source→source transform: full JS in, canonical jz out.
+ * Same jzify/ module the compiler runs on every parse (default-on) — re-exported
+ * here so browser bundles (dist/jz.js: REPL auto-jzify on paste) and node users
+ * (`import { transform } from 'jz'`, also `jz/transform`) share one lowering.
+ */
+export { default as transform } from './transform.js'