jz 0.1.1 → 0.2.0

package/src/compile.js

@@ -26,13 +26,13 @@
  */
 
 import { parse as parseWat } from 'watr'
-import { ctx, err, inc, resolveIncludes, PTR } from './ctx.js'
+import { ctx, err, inc, resolveIncludes, PTR, LAYOUT } from './ctx.js'
 import {
   T, VAL, analyzeValTypes, analyzeIntCertain, analyzeLocals,
   analyzePtrUnboxable, typedElemAux, invalidateLocalsCache,
-  analyzeBoxedCaptures, updateRep,
+  analyzeBoxedCaptures, updateRep, inferStringParams,
 } from './analyze.js'
-import { optimizeFunc, hoistConstantPool, specializeMkptr, specializePtrBase, sortStrPoolByFreq, treeshake } from './optimize.js'
+import { optimizeFunc, hoistConstantPool, specializeMkptr, specializePtrBase, sortStrPoolByFreq, treeshake, arenaRewindModule } from './optimize.js'
 import { emit, emitter, emitFlat, emitBody } from './emit.js'
 import {
   typed, asF64, asI32, asPtrOffset, asParamType, toI32, asI64, fromI64,
@@ -42,7 +42,7 @@ import {
   isLit, litVal, isNullishLit, isPureIR, isPostfix, emitNum,
   temp, tempI32, tempI64, f64rem, toNumF64, truthyIR, toBoolFromEmitted,
   keyValType, usesDynProps, needsDynShadow,
-  isGlobal, isConst, boxedAddr, readVar, writeVar, isNullish,
+  isGlobal, isConst, boxedAddr, readVar, writeVar, isNullish, isUndef,
   slotAddr, elemLoad, elemStore, arrayLoop, allocPtr,
   multiCount, loopTop, flat, reconstructArgsWithSpreads,
   valKindToPtr,
@@ -54,10 +54,15 @@ const timePhase = (profiler, name, fn) => profiler ? profiler.time(name, fn) : f
 // Per-compile func name set + map live on ctx.func.names / ctx.func.map,
 // populated at compile() entry. Both reset by ctx.js reset() and re-filled here.
 
-// NaN-box high-bits mask: used by the static-prefix-strip pass below to
-// identify pointer slots in the data segment. Kept local (ir.js owns the
-// runtime packing via mkPtrIR).
-const NAN_PREFIX_BITS = 0x7FF8n
+// NaN-prefix top-13-bits as BigInt — used by the static-prefix-strip pass
+// below to identify pointer slots in the data segment.
+const NAN_PREFIX = BigInt(LAYOUT.NAN_PREFIX)
+const TAG_MASK_BIG = BigInt(LAYOUT.TAG_MASK)
+const AUX_MASK_BIG = BigInt(LAYOUT.AUX_MASK)
+const OFFSET_MASK_BIG = BigInt(LAYOUT.OFFSET_MASK)
+const TAG_SHIFT_BIG = BigInt(LAYOUT.TAG_SHIFT)
+const AUX_SHIFT_BIG = BigInt(LAYOUT.AUX_SHIFT)
+const SSO_BIT_BIG = BigInt(LAYOUT.SSO_BIT)
 
 // Low-level IR helpers previously lived here. Pure ones moved to src/ir.js;
 // emit-calling ones (toBool, emitTypeofCmp, emitDecl, materializeMulti,
@@ -140,6 +145,95 @@ const tcoTailRewrite = (ir, resultType) => {
   return ir
 }
 
+const ARENA_SAFE_CALLS = new Set([
+  '$__alloc', '$__alloc_hdr', '$__mkptr',
+  '$__ptr_offset', '$__ptr_type', '$__ptr_aux',
+  '$__len', '$__cap', '$__typed_shift', '$__typed_data',
+])
+
+const findFuncBodyStart = (fn) => {
+  for (let i = 2; i < fn.length; i++) {
+    const n = fn[i]
+    if (Array.isArray(n) && (n[0] === 'param' || n[0] === 'result' || n[0] === 'local' || n[0] === 'export')) continue
+    return i
+  }
+  return fn.length
+}
+
+const heapGetIR = () => ctx.memory.shared
+  ? ['i32.load', ['i32.const', 1020]]
+  : ['global.get', '$__heap']
+
+const heapSetIR = value => ctx.memory.shared
+  ? ['i32.store', ['i32.const', 1020], value]
+  : ['global.set', '$__heap', value]
+
+function applyArenaRewind(func, fn, safeCallees) {
+  if (ctx.transform.optimize?.arenaRewind === false) return false
+  if (func.raw || func.sig.params.length !== 0 || func.sig.results.length !== 1) return false
+  if (func.sig.ptrKind != null) return false
+  if (func.sig.results[0] === 'f64' && func.valResult !== VAL.NUMBER) return false
+  if (func.sig.results[0] !== 'f64' && func.sig.results[0] !== 'i32') return false
+
+  const bodyStart = findFuncBodyStart(fn)
+  let hasAlloc = false
+  let unsafe = false
+  const scan = node => {
+    if (unsafe || !Array.isArray(node)) return
+    const op = node[0]
+    if (op === 'global.set' || op === 'return_call' || op === 'call_indirect' || op === 'call_ref') {
+      unsafe = true
+      return
+    }
+    if (op === 'call') {
+      const name = node[1]
+      if (name === '$__alloc' || name === '$__alloc_hdr') hasAlloc = true
+      if (!(safeCallees ?? ARENA_SAFE_CALLS).has(name)) {
+        unsafe = true
+        return
+      }
+    }
+    for (let i = 1; i < node.length; i++) scan(node[i])
+  }
+  for (let i = bodyStart; i < fn.length; i++) scan(fn[i])
+  if (unsafe || !hasAlloc) return false
+
+  let id = 0
+  const hasLocal = name => fn.some(n => Array.isArray(n) && n[0] === 'local' && n[1] === name)
+  while (hasLocal(`$${T}heap_save${id}`) || hasLocal(`$${T}arena_ret${id}`)) id++
+  const save = `$${T}heap_save${id}`
+  const ret = `$${T}arena_ret${id}`
+  const restore = () => heapSetIR(['local.get', save])
+  const resultType = func.sig.results[0]
+
+  const rewriteReturns = node => {
+    if (!Array.isArray(node)) return node
+    if (node[0] === 'return' && node.length > 1) {
+      return ['block',
+        ['result', resultType],
+        ['local.set', ret, node[1]],
+        restore(),
+        ['return', ['local.get', ret]],
+        ['unreachable']]
+    }
+    for (let i = 1; i < node.length; i++) node[i] = rewriteReturns(node[i])
+    return node
+  }
+
+  const endsWithReturn = fn.at(-1)?.[0] === 'return' || fn.at(-1)?.[0] === 'return_call'
+  for (let i = bodyStart; i < fn.length; i++) fn[i] = rewriteReturns(fn[i])
+  const newBodyStart = findFuncBodyStart(fn)
+  fn.splice(newBodyStart, 0,
+    ['local', save, 'i32'],
+    ['local', ret, resultType],
+    ['local.set', save, heapGetIR()])
+  if (!endsWithReturn) {
+    const last = fn.pop()
+    fn.push(['local.set', ret, last], restore(), ['local.get', ret])
+  }
+  return true
+}
+
 // === Module compilation ===
 
 const cloneRepMap = map => map ? new Map([...map].map(([k, v]) => [k, { ...v }])) : null
@@ -186,20 +280,39 @@ function analyzeFuncForEmit(func, programFacts) {
   // cached widths reflect the pre-narrow state. Re-walk now with reps in place.
   invalidateLocalsCache(body)
   ctx.func.locals = block ? analyzeLocals(body) : new Map()
+  // Usage-based VAL.STRING inference for params not already typed by paramReps.
+  // Descends into nested closures so a param used as STRING only inside an inner
+  // arrow (e.g. parseLevel's `str` capture in watr) still gets seeded — the
+  // closure capture path then propagates VAL.STRING via captureValTypes.
+  if (block) {
+    const candidates = sig.params
+      .filter(p => !ctx.func.repByLocal?.get(p.name)?.val)
+      .map(p => p.name)
+    if (candidates.length) {
+      const inferred = inferStringParams(body, candidates)
+      for (const [n, vt] of inferred) updateRep(n, { val: vt })
+    }
+  }
   if (block) {
     analyzeValTypes(body)
     analyzeIntCertain(body)
     analyzeBoxedCaptures(body)
     // Lower provably-monomorphic pointer locals to i32 offset storage.
+    // VAL.TYPED unbox requires a known element ctor (aux byte) — without it,
+    // the use site can't pick the right i32.store{8,16}/i32.store width and
+    // the rebox path can't reconstruct the NaN-box. Heterogeneous decls (two
+    // `let arr = ...` with different ctors, or a multi-ctor ternary) leave
+    // typedElem unset; skip unbox so reads/writes go through `__typed_set_idx`.
     const unbox = analyzePtrUnboxable(body, ctx.func.locals, ctx.func.boxed)
     if (unbox.size > 0) {
       for (const [n, kind] of unbox) {
-        ctx.func.locals.set(n, 'i32')
         const fields = { ptrKind: kind }
         if (kind === VAL.TYPED) {
           const aux = typedElemAux(ctx.types.typedElem?.get(n))
-          if (aux != null) fields.ptrAux = aux
+          if (aux == null) continue
+          fields.ptrAux = aux
         }
+        ctx.func.locals.set(n, 'i32')
         updateRep(n, fields)
       }
     }
@@ -233,7 +346,7 @@ function analyzeFuncForEmit(func, programFacts) {
 function emitFunc(func, funcFacts, programFacts) {
   const { paramReps } = programFacts
 
-  // Raw WAT functions (e.g., _alloc, _reset from memory module)
+  // Raw WAT functions (e.g., _alloc, _clear from memory module)
   if (func.raw) return parseWat(func.raw)
 
   const { name, body, exported, sig } = func
@@ -275,24 +388,28 @@ function emitFunc(func, funcFacts, programFacts) {
   fn.push(...sig.params.map(p => ['param', `$${p.name}`, p.type]))
   fn.push(...sig.results.map(t => ['result', t]))
 
-  // Default params: missing JS args become canonical NaN (0x7FF8000000000000) in WASM f64 params.
-  // Check for canonical NaN specifically — NaN-boxed pointers are also NaN but have non-zero payload.
+  // Default params: ES spec says default applies only when arg is `undefined`
+  // (or missing). `null`, `0`, `false`, etc. all skip the default.
   const defaults = func.defaults || {}
   const defaultInits = []
   for (const [pname, defVal] of Object.entries(defaults)) {
     const p = sig.params.find(p => p.name === pname)
     const t = p?.type || 'f64'
     defaultInits.push(
-      ['if', isNullish(typed(['local.get', `$${pname}`], 'f64')),
+      ['if', isUndef(typed(['local.get', `$${pname}`], 'f64')),
         ['then', ['local.set', `$${pname}`, t === 'f64' ? asF64(emit(defVal)) : asI32(emit(defVal))]]])
   }
 
   // Box params that are mutably captured: allocate cell, copy param value
   const boxedParamInits = []
+  const preboxedLocalInits = []
+  ctx.func.preboxed = new Set()
+  const paramNames = new Set(sig.params.map(p => p.name))
   for (const p of sig.params) {
     if (ctx.func.boxed.has(p.name)) {
       const cell = ctx.func.boxed.get(p.name)
       ctx.func.locals.set(cell, 'i32')
+      ctx.func.preboxed.add(p.name)
       const lget = typed(['local.get', `$${p.name}`], p.type)
       if (p.ptrKind != null) lget.ptrKind = p.ptrKind
       boxedParamInits.push(
@@ -300,6 +417,14 @@ function emitFunc(func, funcFacts, programFacts) {
         ['f64.store', ['local.get', `$${cell}`], asF64(lget)])
     }
   }
+  for (const [name, cell] of ctx.func.boxed) {
+    if (paramNames.has(name)) continue
+    ctx.func.locals.set(cell, 'i32')
+    ctx.func.preboxed.add(name)
+    preboxedLocalInits.push(
+      ['local.set', `$${cell}`, ['call', '$__alloc', ['i32.const', 8]]],
+      ['f64.store', ['local.get', `$${cell}`], nullExpr()])
+  }
 
   if (block) {
     const stmts = emitBody(body)
@@ -307,16 +432,16 @@ function emitFunc(func, funcFacts, programFacts) {
     // I: Skip trailing fallback when last statement is return (unreachable code)
     const lastStmt = stmts.at(-1)
     const endsWithReturn = lastStmt && (lastStmt[0] === 'return' || lastStmt[0] === 'return_call')
-    fn.push(...defaultInits, ...boxedParamInits, ...stmts, ...(endsWithReturn ? [] : sig.results.map(t => [`${t}.const`, 0])))
+    fn.push(...defaultInits, ...boxedParamInits, ...preboxedLocalInits, ...stmts, ...(endsWithReturn ? [] : sig.results.map(t => [`${t}.const`, 0])))
   } else if (multi && body[0] === '[') {
     const values = body.slice(1).map(e => asF64(emit(e)))
     for (const [l, t] of ctx.func.locals) fn.push(['local', `$${l}`, t])
-    fn.push(...boxedParamInits, ...values)
+    fn.push(...boxedParamInits, ...preboxedLocalInits, ...values)
   } else {
     const ir = emit(body)
     for (const [l, t] of ctx.func.locals) fn.push(['local', `$${l}`, t])
     const finalIR = sig.ptrKind != null ? asPtrOffset(ir, sig.ptrKind) : asParamType(ir, sig.results[0])
-    fn.push(...defaultInits, ...boxedParamInits, tcoTailRewrite(finalIR, sig.results[0]))
+    fn.push(...defaultInits, ...boxedParamInits, ...preboxedLocalInits, tcoTailRewrite(finalIR, sig.results[0]))
   }
 
   // Restore schema.vars so param bindings don't leak to next function.
@@ -329,18 +454,22 @@ function emitFunc(func, funcFacts, programFacts) {
  *
  * For each `isBoundaryWrapped(func)`, emit a sibling `$${name}$exp` that:
  *   - holds the (export "name") attribute (JS sees the wrapper)
- *   - takes f64 params always (JS calling convention via host.js wrap)
+ *   - takes i64 params always — JS-side carrier is BigInt that reinterprets to
+ *     f64 NaN-box bits. i64 dodges V8's spec-permitted NaN canonicalization at
+ *     the wasm↔JS boundary (see ToJSValue / ToWebAssemblyValue). Host wrap()
+ *     in src/host.js pairs by converting BigInt↔f64 via reinterpret bits.
  *   - converts each narrowed param at the call: f64 → i32 (truncate-sat) for
  *     numeric narrowed, f64 → i32-offset (`i32.wrap_i64 + i64.reinterpret_f64`)
- *     for pointer narrowed
+ *     for pointer narrowed. The reinterpret happens once at param decode and
+ *     once at result encode; numeric exports without narrowing skip wrapping
+ *     entirely (no NaN-class values).
  *   - forwards args to the inner $${name}
- *   - reboxes the narrowed result back to f64 so JS sees Number / NaN-boxed ptr
+ *   - reboxes the narrowed result and reinterprets to i64 for the boundary
  *
- * Param convert cases (each narrowed inner-param):
- *   - p.type = 'i32', no ptrKind  → i32.trunc_sat_f64_s(local.get $p)
- *   - p.type = 'i32', ptrKind set → i32.wrap_i64(i64.reinterpret_f64(local.get $p))
+ * Param decode (i64 → f64): each param gets `f64.reinterpret_i64` before the
+ * existing narrowing convert. f64 inner params just need the reinterpret.
  *
- * Result rebox cases:
+ * Result rebox cases (then reinterpret to i64 at the boundary):
  *   - sig.ptrKind != null  → mkPtrIR(ptrKind, ptrAux ?? 0, callIR)
  *   - sig.results[0] = i32 → f64.convert_i32_s(callIR)
  *   - sig.results[0] = f64 → callIR (some params narrowed but result stayed f64)
@@ -350,18 +479,23 @@ function synthesizeBoundaryWrappers() {
   for (const func of ctx.func.list) {
     if (!isBoundaryWrapped(func)) continue
     const { name, sig } = func
+    // Per-position i64 carrier: only swap to i64 where a NaN-boxed pointer
+    // actually crosses the boundary (param.ptrKind set, or result with
+    // sig.ptrKind set). Numeric narrowing (i32 trunc-sat / convert) keeps f64
+    // so callers seeing the raw export get a plain Number for numerics.
+    const paramI64 = sig.params.map(p => p.ptrKind != null)
+    const resultI64 = sig.ptrKind != null
     const wrapNode = ['func', `$${name}$exp`, ['export', `"${name}"`]]
-    // External ABI: every param is f64 (JS Number / NaN-boxed ptr).
-    for (const p of sig.params) wrapNode.push(['param', `$${p.name}`, 'f64'])
-    wrapNode.push(['result', 'f64'])
-    const args = sig.params.map(p => {
+    sig.params.forEach((p, i) => wrapNode.push(['param', `$${p.name}`, paramI64[i] ? 'i64' : 'f64']))
+    wrapNode.push(['result', resultI64 ? 'i64' : 'f64'])
+    const args = sig.params.map((p, i) => {
       const get = ['local.get', `$${p.name}`]
-      if (p.type === 'f64') return get
       if (p.ptrKind != null) {
-        // NaN-boxed f64 → raw i32 offset
-        return ['i32.wrap_i64', ['i64.reinterpret_f64', get]]
+        // ptrKind: i64 carrier carries NaN-box bits → wrap to i32 offset
+        return ['i32.wrap_i64', get]
       }
-      // Numeric i32 — JS Number → i32 truncation (matches `n | 0` for integers).
+      if (p.type === 'f64') return get
+      // Numeric narrowing: f64 → i32 truncate
       return ['i32.trunc_sat_f64_s', get]
     })
     const callIR = ['call', `$${name}`, ...args]
@@ -374,7 +508,9 @@ function synthesizeBoundaryWrappers() {
     } else {
       body = callIR
     }
-    wrapNode.push(body)
+    wrapNode.push(resultI64 ? ['i64.reinterpret_f64', body] : body)
+    func._exportUsesI64 = resultI64 || paramI64.some(Boolean)
+    func._exportI64Sig = { params: paramI64, result: resultI64 }
     wrappers.push(wrapNode)
   }
   return wrappers
@@ -398,6 +534,7 @@ function emitClosureBody(cb) {
   // Reset per-function state for closure body
   ctx.func.locals = new Map()
   ctx.func.repByLocal = null
+  if (cb.intConsts) for (const [name, v] of cb.intConsts) updateRep(name, { intConst: v })
   if (cb.valTypes) for (const [name, vt] of cb.valTypes) updateRep(name, { val: vt })
   if (cb.schemaVars) {
     ctx.schema.vars = new Map([...prevSchemaVars, ...cb.schemaVars])
@@ -413,6 +550,8 @@ function emitClosureBody(cb) {
   }
   // In closure bodies, boxed captures use the original name as both var and cell local
   ctx.func.boxed = cb.boxed ? new Map([...cb.boxed].map(v => [v, v])) : new Map()
+  const parentBoxedCaptures = new Set(cb.boxed || [])
+  ctx.func.preboxed = new Set()
   ctx.func.stack = []
   ctx.func.uniq = Math.max(ctx.func.uniq, 100) // avoid label collisions
   ctx.func.body = cb.body
@@ -445,11 +584,35 @@ function emitClosureBody(cb) {
   const block = Array.isArray(cb.body) && cb.body[0] === '{}' && cb.body[1]?.[0] !== ':'
   let bodyIR
   if (block) {
+    invalidateLocalsCache(cb.body)
     for (const [k, v] of analyzeLocals(cb.body)) if (!ctx.func.locals.has(k)) ctx.func.locals.set(k, v)
+    // Usage-based STRING inference for closure params not seeded by captureValTypes.
+    // (Captures already have their parent's val type via cb.valTypes above.)
+    {
+      const candidates = cb.params.filter(p => !ctx.func.repByLocal?.get(p)?.val)
+      if (candidates.length) {
+        const inferred = inferStringParams(cb.body, candidates)
+        for (const [n, vt] of inferred) updateRep(n, { val: vt })
+      }
+    }
+    analyzeValTypes(cb.body)
+    analyzeIntCertain(cb.body)
     // Detect captures from deeper nested arrows that mutate this body's locals/params/captures
     analyzeBoxedCaptures(cb.body)
     for (const name of ctx.func.boxed.keys()) {
-      if (ctx.func.locals.get(name) === 'f64') ctx.func.locals.set(name, 'i32')
+      if (parentBoxedCaptures.has(name) && ctx.func.locals.get(name) === 'f64') ctx.func.locals.set(name, 'i32')
+    }
+    const unbox = analyzePtrUnboxable(cb.body, ctx.func.locals, ctx.func.boxed)
+    for (const [name, kind] of unbox) {
+      if (cb.params.includes(name) || cb.captures.includes(name)) continue
+      const fields = { ptrKind: kind }
+      if (kind === VAL.TYPED) {
+        const aux = typedElemAux(ctx.types.typedElem?.get(name))
+        if (aux == null) continue
+        fields.ptrAux = aux
+      }
+      ctx.func.locals.set(name, 'i32')
+      updateRep(name, fields)
     }
     bodyIR = emitBody(cb.body)
   } else {
@@ -469,6 +632,28 @@ function emitClosureBody(cb) {
     inc('__alloc_hdr', '__mkptr')
   }
 
+  const boxedCaptureNames = new Set(cb.captures.filter(name => parentBoxedCaptures.has(name)))
+  for (const name of boxedCaptureNames) ctx.func.preboxed.add(name)
+  const boxedValueCaptureNames = new Set(cb.captures.filter(name => ctx.func.boxed.has(name) && !parentBoxedCaptures.has(name)))
+  for (const name of boxedValueCaptureNames) {
+    ctx.func.locals.set(ctx.func.boxed.get(name), 'i32')
+    ctx.func.preboxed.add(name)
+  }
+  const boxedParamNames = new Set(cb.params.filter(name => ctx.func.boxed.has(name)))
+  for (const name of boxedParamNames) {
+    ctx.func.locals.set(ctx.func.boxed.get(name), 'i32')
+    ctx.func.preboxed.add(name)
+  }
+  const preboxedLocalInits = []
+  for (const [name, cell] of ctx.func.boxed) {
+    if (boxedCaptureNames.has(name) || boxedValueCaptureNames.has(name) || boxedParamNames.has(name)) continue
+    ctx.func.locals.set(cell, 'i32')
+    ctx.func.preboxed.add(name)
+    preboxedLocalInits.push(
+      ['local.set', `$${cell}`, ['call', '$__alloc', ['i32.const', 8]]],
+      ['f64.store', ['local.get', `$${cell}`], nullExpr()])
+  }
+
   // Insert locals (captures + params + declared)
   for (const [l, t] of ctx.func.locals) fn.push(['local', `$${l}`, t])
 
@@ -481,8 +666,15 @@ function emitClosureBody(cb) {
     for (let i = 0; i < cb.captures.length; i++) {
       const name = cb.captures[i]
       const addr = ['i32.add', ['local.get', `$${envBase}`], ['i32.const', i * 8]]
-      fn.push(['local.set', `$${name}`,
-        ctx.func.boxed.has(name) ? ['i32.load', addr] : ['f64.load', addr]])
+      if (parentBoxedCaptures.has(name)) {
+        fn.push(['local.set', `$${name}`, ['i32.load', addr]])
+      } else if (boxedValueCaptureNames.has(name)) {
+        fn.push(
+          ['local.set', `$${ctx.func.boxed.get(name)}`, ['call', '$__alloc', ['i32.const', 8]]],
+          ['f64.store', boxedAddr(name), ['f64.load', addr]])
+      } else {
+        fn.push(['local.set', `$${name}`, ['f64.load', addr]])
+      }
     }
   }
 
@@ -490,7 +682,14 @@ function emitClosureBody(cb) {
   // Rest name (if present) is last in cb.params — handled separately below.
   const fixedParamN = cb.params.length - (cb.rest ? 1 : 0)
   for (let i = 0; i < fixedParamN && i < W; i++) {
-    fn.push(['local.set', `$${cb.params[i]}`, ['local.get', `$__a${i}`]])
+    const pname = cb.params[i]
+    if (boxedParamNames.has(pname)) {
+      fn.push(
+        ['local.set', `$${ctx.func.boxed.get(pname)}`, ['call', '$__alloc', ['i32.const', 8]]],
+        ['f64.store', boxedAddr(pname), ['local.get', `$__a${i}`]])
+    } else {
+      fn.push(['local.set', `$${pname}`, ['local.get', `$__a${i}`]])
+    }
   }
 
   // Rest param: pack slots a[fixedParams..argc-1] into fresh array.
@@ -516,20 +715,34 @@ function emitClosureBody(cb) {
           ['i32.add', ['local.get', `$${restOff}`], ['i32.const', i * 8]],
           ['local.get', `$__a${fixedN + i}`]]]])
     }
-    fn.push(['local.set', `$${cb.rest}`,
-      ['call', '$__mkptr', ['i32.const', PTR.ARRAY], ['i32.const', 0], ['local.get', `$${restOff}`]]])
+    const restValue = ['call', '$__mkptr', ['i32.const', PTR.ARRAY], ['i32.const', 0], ['local.get', `$${restOff}`]]
+    if (boxedParamNames.has(cb.rest)) {
+      fn.push(
+        ['local.set', `$${ctx.func.boxed.get(cb.rest)}`, ['call', '$__alloc', ['i32.const', 8]]],
+        ['f64.store', boxedAddr(cb.rest), restValue])
+    } else {
+      fn.push(['local.set', `$${cb.rest}`, restValue])
+    }
   }
 
   // Default params for closures (check sentinel after unpack)
+  // Only `undefined` triggers default per spec — `null`/`0`/`false` pass through.
   if (cb.defaults) {
     for (const [pname, defVal] of Object.entries(cb.defaults)) {
-      fn.push(['if', isNullish(['local.get', `$${pname}`]),
-        ['then', ['local.set', `$${pname}`, asF64(emit(defVal))]]])
+      if (boxedParamNames.has(pname)) {
+        fn.push(['if', isUndef(['f64.load', boxedAddr(pname)]),
+          ['then', ['f64.store', boxedAddr(pname), asF64(emit(defVal))]]])
+      } else {
+        fn.push(['if', isUndef(['local.get', `$${pname}`]),
+          ['then', ['local.set', `$${pname}`, asF64(emit(defVal))]]])
+      }
     }
   }
+  fn.push(...preboxedLocalInits)
   fn.push(...bodyIR)
   // I: Skip trailing fallback when last statement is return
-  if (block && !(bodyIR.at(-1)?.[0] === 'return' || bodyIR.at(-1)?.[0] === 'return_call')) fn.push(['f64.const', 0])
+  // Implicit fall-through return is `undefined` per JS spec, not 0.
+  if (block && !(bodyIR.at(-1)?.[0] === 'return' || bodyIR.at(-1)?.[0] === 'return_call')) fn.push(undefExpr())
   ctx.schema.vars = prevSchemaVars
   ctx.types.typedElem = prevTypedElems
   return fn
@@ -582,9 +795,9 @@ function buildStartFn(ast, sec, closureFuncs, compilePendingClosures) {
     const bt = `${T}box`
     ctx.func.locals.set(bt, 'i32')
     for (const [name, { schemaId, schema }] of ctx.schema.autoBox) {
-      inc('__alloc', '__mkptr')
+      inc('__alloc_hdr', '__mkptr')
       boxInit.push(
-        ['local.set', `$${bt}`, ['call', '$__alloc', ['i32.const', schema.length * 8]]],
+        ['local.set', `$${bt}`, ['call', '$__alloc_hdr', ['i32.const', 0], ['i32.const', Math.max(1, schema.length)], ['i32.const', 8]]],
         ['f64.store', ['local.get', `$${bt}`],
           ctx.func.names.has(name) ? ['f64.const', 0] : ['global.get', `$${name}`]],
         ...schema.slice(1).map((_, i) =>
@@ -602,7 +815,12 @@ function buildStartFn(ast, sec, closureFuncs, compilePendingClosures) {
   // __dyn_get (set transitively by resolveIncludes() later) are listed
   // explicitly here because the dep graph hasn't been expanded yet at
   // start-fn build time.
-  const needsSchemaTbl = ctx.schema.list.length && (
+  // __jp_obj registers schemas at runtime, so the table must be allocated
+  // (and __schema_next initialized) even with zero compile-time schemas.
+  // __jp pulls __jp_obj transitively via the include-graph but resolveIncludes
+  // runs after buildStartFn, so check the top-level __jp parser as a proxy.
+  const hasJpObj = ctx.core.includes.has('__jp_obj') || ctx.core.includes.has('__jp')
+  const needsSchemaTbl = (ctx.schema.list.length && (
     ctx.core.includes.has('__stringify') ||
     ctx.core.includes.has('__dyn_get') ||
     ctx.core.includes.has('__dyn_get_t') ||
@@ -610,17 +828,24 @@ function buildStartFn(ast, sec, closureFuncs, compilePendingClosures) {
     ctx.core.includes.has('__dyn_get_any_t') ||
     ctx.core.includes.has('__dyn_get_expr') ||
     ctx.core.includes.has('__dyn_get_expr_t') ||
-    ctx.core.includes.has('__dyn_get_or'))
+    ctx.core.includes.has('__dyn_get_or'))) ||
+    hasJpObj
   if (needsSchemaTbl) {
     const nSchemas = ctx.schema.list.length
+    // Reserve trailing slots for runtime-registered schemas (used by JSON.parse
+    // shape caching). __schema_next tracks the first free runtime slot.
+    const runtimeReserve = hasJpObj ? 256 : 0
     const stbl = `${T}stbl`
     const sarr = `${T}sarr`
     ctx.func.locals.set(stbl, 'i32')
     ctx.func.locals.set(sarr, 'i32')
     inc('__alloc', '__alloc_hdr', '__mkptr')
     schemaInit.push(
-      ['local.set', `$${stbl}`, ['call', '$__alloc', ['i32.const', nSchemas * 8]]],
+      ['local.set', `$${stbl}`, ['call', '$__alloc', ['i32.const', (nSchemas + runtimeReserve) * 8]]],
       ['global.set', '$__schema_tbl', ['local.get', `$${stbl}`]])
+    if (runtimeReserve) {
+      schemaInit.push(['global.set', '$__schema_next', ['i32.const', nSchemas]])
+    }
     for (let s = 0; s < nSchemas; s++) {
       const keys = ctx.schema.list[s]
       const n = keys.length
@@ -651,12 +876,15 @@ function buildStartFn(ast, sec, closureFuncs, compilePendingClosures) {
     for (const s of ctx.runtime.typeofStrs)
       typeofInit.push(['global.set', `$__tof_${s}`, emit(['str', s])])
   }
-  if (moduleInits.length || init?.length || boxInit.length || schemaInit.length || typeofInit.length || strPoolInit.length || ctx.features.timers) {
+  // WASI timer queue needs an init call in __start; JS-host mode delegates
+  // scheduling to the host so no init/loop is needed.
+  const wasiTimers = ctx.features.timers && ctx.transform.host === 'wasi'
+  if (moduleInits.length || init?.length || boxInit.length || schemaInit.length || typeofInit.length || strPoolInit.length || wasiTimers) {
     const initIR = normalizeIR(init)
     const startFn = ['func', '$__start']
     for (const [l, t] of ctx.func.locals) startFn.push(['local', `$${l}`, t])
     startFn.push(...strPoolInit, ...typeofInit, ...boxInit, ...schemaInit,
-      ...(ctx.features.timers ? [['call', '$__timer_init']] : []),
+      ...(wasiTimers ? [['call', '$__timer_init']] : []),
       ...moduleInits, ...initIR,
       ...(ctx.features.blockingTimers ? [['call', '$__timer_loop']] : []),
     )
@@ -721,12 +949,21 @@ function dedupClosureBodies(closureFuncs, sec) {
     else { hashToName.set(key, name); keepSet.add(name) }
   }
   if (!redirect.size) return
-  ctx.closure.table = ctx.closure.table.map(n => redirect.get(n) || n)
+  // Remove duplicate declarations BEFORE rewriting references — otherwise redirectRefs
+  // renames the declaration itself and the filter can't find the original name.
   const kept = sec.funcs.filter(fn => {
     if (!Array.isArray(fn) || fn[0] !== 'func') return true
     const name = typeof fn[1] === 'string' && fn[1][0] === '$' ? fn[1].slice(1) : null
     return !name || !redirect.has(name)
   })
+  const redirectRefs = node => {
+    if (typeof node === 'string') return node[0] === '$' && redirect.has(node.slice(1)) ? `$${redirect.get(node.slice(1))}` : node
+    if (!Array.isArray(node)) return node
+    for (let i = 0; i < node.length; i++) node[i] = redirectRefs(node[i])
+    return node
+  }
+  for (const fn of kept) redirectRefs(fn)
+  ctx.closure.table = ctx.closure.table.map(n => redirect.get(n) || n)
   sec.funcs.length = 0
   sec.funcs.push(...kept)
 }
@@ -752,8 +989,7 @@ function dedupClosureBodies(closureFuncs, sec) {
  *      Both `call` and `return_call` (tail call) sites are rewritten in the same walk.
  */
 function finalizeClosureTable(sec) {
-  if (!ctx.closure.table?.length) return
-  let indirectUsed = false
+  let indirectUsed = ctx.transform.host === 'wasi'
   const scan = (n) => {
     if (!Array.isArray(n) || indirectUsed) return
     if (n[0] === 'call_indirect') { indirectUsed = true; return }
@@ -767,8 +1003,9 @@ function finalizeClosureTable(sec) {
   }
   // Keep table if call_indirect is used (closures, timer dispatch, etc.)
   if (indirectUsed) {
+    if (!ctx.closure.table) ctx.closure.table = []
     sec.table = [['table', ['export', '"__jz_table"'], ctx.closure.table.length, 'funcref']]
-    sec.elem = [['elem', ['i32.const', 0], 'func', ...ctx.closure.table.map(n => `$${n}`)]]
+    sec.elem = ctx.closure.table.length ? [['elem', ['i32.const', 0], 'func', ...ctx.closure.table.map(n => `$${n}`)]] : []
     return
   }
   sec.table = []
@@ -832,7 +1069,7 @@ function finalizeClosureTable(sec) {
  *   1. resolveIncludes() — close the include set under stdlib dependencies.
  *   2. Emit memory section ONLY when some included helper uses memory ops
  *      (G optimization: pure scalar programs ship without memory + __heap).
- *      When memory is needed, the allocator (__alloc + __alloc_hdr + __reset)
+ *      When memory is needed, the allocator (__alloc + __alloc_hdr + __clear)
  *      is force-included since stdlib funcs may call into it.
  *   3. Pull external (host) stdlibs into sec.extStdlib (must precede normal
  *      imports in the module byte order).
@@ -848,11 +1085,11 @@ function pullStdlib(sec) {
   const needsMemory = [...ctx.core.includes].some(n => ctx.core.stdlib[n] && MEM_OPS.test(ctx.core.stdlib[n]))
   if (!needsMemory) ctx.scope.globals.delete('__heap')
   if (needsMemory && ctx.module.modules.core) {
-    for (const fn of ['__alloc', '__alloc_hdr', '__reset']) if (!ctx.core.includes.has(fn)) ctx.core.includes.add(fn)
+    for (const fn of ['__alloc', '__alloc_hdr', '__clear']) if (!ctx.core.includes.has(fn)) ctx.core.includes.add(fn)
     const pages = ctx.memory.pages || 1
     if (ctx.memory.shared) sec.imports.push(['import', '"env"', '"memory"', ['memory', pages]])
     else sec.memory.push(['memory', ['export', '"memory"'], pages])
-    if (ctx.transform.runtimeExports !== false && ctx.core._allocRawFuncs)
+    if (ctx.transform.alloc !== false && ctx.core._allocRawFuncs)
       sec.funcs.push(...ctx.core._allocRawFuncs.map(s => parseWat(s)))
   }
 
@@ -860,10 +1097,15 @@ function pullStdlib(sec) {
     const v = ctx.core.stdlib[name]
     return typeof v === 'function' ? v() : v
   }
+  // Track __ext_* names emitted to host imports. The cleanup below deletes them
+  // from ctx.core.includes (moved into sec.extStdlib instead), so any post-compile
+  // audit (e.g. host: 'wasi') needs to read this list rather than the includes set.
+  ctx.core.extImports ??= new Set()
   for (const name of Object.keys(ctx.core.stdlib)) {
     if (name.startsWith('__ext_') && ctx.core.includes.has(name)) {
       const parsed = parseWat(stdlibStr(name))
       sec.extStdlib.push(parsed[0] === "module" ? parsed[1] : parsed)
+      ctx.core.extImports.add(name)
       ctx.core.includes.delete(name)
     }
   }
@@ -871,6 +1113,12 @@ function pullStdlib(sec) {
   sec.stdlib.push(...[...ctx.core.includes].map(n => parseWat(stdlibStr(n))))
 }
 
+function syncImports(sec) {
+  for (const imp of ctx.module.imports) {
+    if (!sec.imports.some(i => i[1] === imp[1] && i[2] === imp[2])) sec.imports.push(imp)
+  }
+}
+
 /**
  * Phase: whole-module + per-function optimization passes.
  *
@@ -904,6 +1152,19 @@ function optimizeModule(sec) {
     ctx.runtime.strPool = poolRef.pool
   }
   for (const s of [...sec.funcs, ...sec.stdlib, ...sec.start]) optimizeFunc(s, cfg)
+  if (!cfg || cfg.arenaRewind !== false) {
+    const safeCallees = arenaRewindModule([...sec.funcs, ...sec.stdlib, ...sec.start])
+    // Build name → WAT IR map for user functions
+    const fnByName = new Map()
+    for (const fn of sec.funcs) {
+      if (Array.isArray(fn) && fn[0] === 'func' && typeof fn[1] === 'string')
+        fnByName.set(fn[1], fn)
+    }
+    for (const func of ctx.func.list) {
+      const fn = fnByName.get(`$${func.name}`)
+      if (fn) applyArenaRewind(func, fn, safeCallees)
+    }
+  }
   if (!cfg || cfg.hoistConstantPool !== false)
     hoistConstantPool([...sec.funcs, ...sec.stdlib, ...sec.start], (name, wat) => ctx.scope.globals.set(name, wat))
 
@@ -911,8 +1172,10 @@ function optimizeModule(sec) {
   if (dataLen > 1024 && !ctx.memory.shared) {
     const heapBase = (dataLen + 7) & ~7
     ctx.scope.globals.set('__heap', `(global $__heap (mut i32) (i32.const ${heapBase}))`)
+    if (ctx.scope.globals.has('__heap_start'))
+      ctx.scope.globals.set('__heap_start', `(global $__heap_start (mut i32) (i32.const ${heapBase}))`)
     for (const s of sec.stdlib)
-      if (s[0] === 'func' && s[1] === '$__reset')
+      if (s[0] === 'func' && s[1] === '$__clear')
         for (let i = 2; i < s.length; i++)
           if (Array.isArray(s[i]) && s[i][0] === 'global.set' && Array.isArray(s[i][2]) && s[i][2][0] === 'i32.const')
             s[i][2][1] = `${heapBase}`
@@ -942,12 +1205,14 @@ function stripStaticDataPrefix(sec) {
     for (const slotOff of ctx.runtime.staticPtrSlots) {
       if (slotOff < prefix) continue
       const bits = dv.getBigUint64(slotOff, true)
-      if (((bits >> 48n) & 0xFFF8n) !== NAN_PREFIX_BITS) continue
-      const ty = Number((bits >> 47n) & 0xFn)
+      if (((bits >> 48n) & 0xFFF8n) !== NAN_PREFIX) continue
+      const ty = Number((bits >> TAG_SHIFT_BIG) & TAG_MASK_BIG)
       if (!SHIFTABLE.has(ty)) continue
-      const off = Number(bits & 0xFFFFFFFFn)
+      // SSO STRING: "offset" holds packed bytes, not a heap address — never shift.
+      if (ty === PTR.STRING && ((bits >> AUX_SHIFT_BIG) & SSO_BIT_BIG)) continue
+      const off = Number(bits & OFFSET_MASK_BIG)
       if (off < prefix) continue
-      const hi = bits & ~0xFFFFFFFFn
+      const hi = bits & ~OFFSET_MASK_BIG
       dv.setBigUint64(slotOff, hi | BigInt(off - prefix), true)
     }
   }
@@ -965,16 +1230,21 @@ function stripStaticDataPrefix(sec) {
         Array.isArray(child[2]) && SHIFTABLE.has(child[2][1]) &&
         Array.isArray(child[4]) && child[4][0] === 'i32.const' &&
         typeof child[4][1] === 'number' && child[4][1] >= prefix) {
-        child[4][1] -= prefix
+        // SSO STRING: aux carries SSO_BIT, "offset" holds packed bytes — don't shift.
+        const isSsoString = child[2][1] === PTR.STRING &&
+          Array.isArray(child[3]) && child[3][0] === 'i32.const' &&
+          typeof child[3][1] === 'number' && (child[3][1] & LAYOUT.SSO_BIT)
+        if (!isSsoString) child[4][1] -= prefix
       } else if (child[0] === 'f64.const' &&
         typeof child[1] === 'string' && child[1].startsWith('nan:0x')) {
         const bits = BigInt(child[1].slice(4)) | 0x7FF0000000000000n
-        if (((bits >> 48n) & 0xFFF8n) === NAN_PREFIX_BITS) {
-          const ty = Number((bits >> 47n) & 0xFn)
-          if (SHIFTABLE.has(ty)) {
-            const off = Number(bits & 0xFFFFFFFFn)
+        if (((bits >> 48n) & 0xFFF8n) === NAN_PREFIX) {
+          const ty = Number((bits >> TAG_SHIFT_BIG) & TAG_MASK_BIG)
+          if (SHIFTABLE.has(ty) &&
+              !(ty === PTR.STRING && ((bits >> AUX_SHIFT_BIG) & SSO_BIT_BIG))) {
+            const off = Number(bits & OFFSET_MASK_BIG)
             if (off >= prefix) {
-              const hi = bits & ~0xFFFFFFFFn
+              const hi = bits & ~OFFSET_MASK_BIG
               const newBits = hi | BigInt(off - prefix)
               child[1] = 'nan:0x' + newBits.toString(16).toUpperCase().padStart(16, '0')
             }
@@ -997,9 +1267,24 @@ export default function compile(ast, profiler) {
   ctx.func.names.clear()
   ctx.func.map.clear()
   for (const f of ctx.func.list) { ctx.func.names.add(f.name); ctx.func.map.set(f.name, f) }
-  // Include imported functions for call resolution (e.g. template interpolations)
-  for (const imp of ctx.module.imports)
-    if (imp[3]?.[0] === 'func') ctx.func.names.add(imp[3][1].replace(/^\$/, ''))
+  // Include imported functions for call resolution (e.g. template interpolations).
+  // Also register a synthesized sig in func.map so emit's arity-aware branches see
+  // the import's declared param count — needed for arg pad/truncate to match it.
+  for (const imp of ctx.module.imports) {
+    if (imp[3]?.[0] !== 'func') continue
+    const fname = imp[3][1].replace(/^\$/, '')
+    ctx.func.names.add(fname)
+    if (!ctx.func.map.has(fname)) {
+      const params = []
+      let result = 'f64'
+      for (let k = 2; k < imp[3].length; k++) {
+        const part = imp[3][k]
+        if (Array.isArray(part) && part[0] === 'param') params.push({ type: part[1] || 'f64' })
+        else if (Array.isArray(part) && part[0] === 'result') result = part[1] || 'f64'
+      }
+      ctx.func.map.set(fname, { name: fname, sig: { params, results: [result] } })
+    }
+  }
 
   // Check user globals don't conflict with runtime globals (modules loaded after user decls)
   for (const name of ctx.scope.userGlobals)
@@ -1106,11 +1391,33 @@ export default function compile(ast, profiler) {
 
   sec.funcs.push(...closureFuncs, ...funcs)
 
+  // WASI command-mode entries (`run`, `_start`) must export as () -> ();
+  // wasmtime/wasmer reject f64-returning functions under those names.
+  // Parametric entries skip this — a CLI invocation has no way to supply args.
+  if (ctx.transform.host === 'wasi') {
+    const WASI_ENTRIES = new Set(['run', '_start'])
+    for (const func of ctx.func.list) {
+      if (!func.exported || !WASI_ENTRIES.has(func.name)) continue
+      if (func.sig.params.length) continue
+      const inner = isBoundaryWrapped(func) ? `$${func.name}$exp` : `$${func.name}`
+      for (const f of sec.funcs) {
+        if (f[1] === inner || f[1] === `$${func.name}`) {
+          const expIdx = f.findIndex(n => Array.isArray(n) && n[0] === 'export')
+          if (expIdx >= 0) f.splice(expIdx, 1)
+        }
+      }
+      sec.funcs.push(['func', `$${func.name}$wasi`, ['export', `"${func.name}"`],
+        ['drop', ['call', inner]]])
+    }
+  }
+
   if (ctx.closure.table?.length)
     sec.elem.push(['elem', ['i32.const', 0], 'func', ...ctx.closure.table.map(n => `$${n}`)])
 
   buildStartFn(ast, sec, closureFuncs, compilePendingClosures)
 
+  syncImports(sec)
+
   dedupClosureBodies(closureFuncs, sec)
 
   finalizeClosureTable(sec)
@@ -1164,6 +1471,28 @@ export default function compile(ast, profiler) {
   if (restParamFuncs.length)
     sec.customs.push(['@custom', '"jz:rest"', `"${JSON.stringify(restParamFuncs).replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"`])
 
+  // Custom section: per-export i64 ABI map. Each entry describes an export
+  // whose boundary wrapper carries NaN-boxed pointers via i64 (rather than
+  // f64) to dodge V8's NaN canonicalization. Format: { name, p, r } where p
+  // is an array of i64 param indices and r is 1 if result is i64. host.js
+  // wrap() reinterprets BigInt↔f64 at i64 positions; numeric f64 positions
+  // stay as Numbers on the JS side.
+  const i64Exports = []
+  for (const f of ctx.func.list) {
+    if (!f.exported || !isBoundaryWrapped(f) || !f._exportUsesI64) continue
+    const p = []
+    f._exportI64Sig.params.forEach((b, i) => { if (b) p.push(i) })
+    const r = f._exportI64Sig.result ? 1 : 0
+    i64Exports.push({ name: f.name, p, r })
+    // Aliases (export { foo as bar }) re-export the same wrapper under a
+    // different JS-visible name; list each alias too so wrap() finds it.
+    for (const [alias, val] of Object.entries(ctx.func.exports)) {
+      if (val === f.name && alias !== f.name) i64Exports.push({ name: alias, p, r })
+    }
+  }
+  if (i64Exports.length)
+    sec.customs.push(['@custom', '"jz:i64exp"', `"${JSON.stringify(i64Exports).replace(/\\/g, '\\\\').replace(/"/g, '\\"')}"`])
+
   // Named export aliases: export { name } or export { source as alias }
   for (const [name, val] of Object.entries(ctx.func.exports)) {
     if (val === true) {