jz 0.1.0 → 0.2.0

package/src/optimize.js

@@ -27,8 +27,13 @@
  * @module optimize
  */
 
+import { LAYOUT } from './ctx.js'
+import { vectorizeLaneLocal } from './vectorize.js'
+
 const MEMOP = /^[fi](32|64)\.(load|store)(\d+(_[su])?)?$/
-const NAN_PREFIX_BITS = 0x7FF8n
+const NAN_BITS = '0x' + LAYOUT.NAN_PREFIX_BITS.toString(16).toUpperCase().padStart(16, '0')
+const NULL_BITS = '0x' + (LAYOUT.NAN_PREFIX_BITS | (1n << BigInt(LAYOUT.AUX_SHIFT))).toString(16).toUpperCase().padStart(16, '0')
+const UNDEF_BITS = '0x' + (LAYOUT.NAN_PREFIX_BITS | (2n << BigInt(LAYOUT.AUX_SHIFT))).toString(16).toUpperCase().padStart(16, '0')
 
 /**
  * Optimization passes, partitioned by phase. The `level` presets pick which
@@ -39,23 +44,30 @@ const NAN_PREFIX_BITS = 0x7FF8n
  *   0 — nothing. Fastest compile, largest output. Useful for live coding.
  *   1 — encoding-compactness only (treeshake + sortLocalsByUse + fusedRewrite-inline).
  *       Cheap, no IR rewrites that perturb V8's tier-up shape.
- *   2 — default. All current passes (watr CSE/DCE/inline + every jz pass).
- *   3 — reserved for future aggressive passes (constant-arg propagation,
- *       inlining, unrolling). Currently == level 2.
+ *   2 — default. Stable passes (watr CSE/DCE/inline + every jz pass that has
+ *       benchmark proof).
+ *   3 — aggressive experimental passes in addition to level 2.
  */
 export const PASS_NAMES = [
   'watr',                     // third-party WAT-level CSE/DCE/inlining (heaviest)
   'hoistPtrType',
   'hoistInvariantPtrOffset',
+  'hoistInvariantPtrOffsetLoop',
   'fusedRewrite',             // peephole + ptr-helper inline + memarg fold
   'hoistAddrBase',
   'hoistInvariantCellLoads',
+  'cseScalarLoad',
+  'csePureExpr',
   'sortLocalsByUse',
   'specializeMkptr',
   'specializePtrBase',
   'sortStrPoolByFreq',
   'hoistConstantPool',
+  'sourceInline',
   'smallConstForUnroll',
+  'nestedSmallConstForUnroll',
+  'vectorizeLaneLocal',       // SIMD-128 lift for lane-pure typed-array loops
+  'arenaRewind',              // per-call heap rewind for no-arg scalar allocator kernels
   'treeshake',
 ]
 
@@ -64,14 +76,14 @@ const ALL_OFF = Object.freeze(Object.fromEntries(PASS_NAMES.map(n => [n, false])
 const LEVEL_PRESETS = Object.freeze({
   0: ALL_OFF,
   1: Object.freeze({ ...ALL_OFF, treeshake: true, sortLocalsByUse: true, fusedRewrite: true }),
-  2: ALL_ON,
+  2: Object.freeze({ ...ALL_ON, nestedSmallConstForUnroll: 'auto' }),
   3: ALL_ON,
 })
 
 /**
  * Normalize the user's `opts.optimize` value into a flat config object.
  *
- *   resolveOptimize(undefined | true)         → all on (level 2)
+ *   resolveOptimize(undefined | true)         → level 2 stable defaults
  *   resolveOptimize(false | 0)                → all off
  *   resolveOptimize(1 | 2 | 3)                → preset for that level
  *   resolveOptimize({ level: 1, watr: true }) → level 1 base, with watr forced on
@@ -79,13 +91,13 @@ const LEVEL_PRESETS = Object.freeze({
  */
 export function resolveOptimize(opt) {
   if (opt === false || opt === 0) return { ...ALL_OFF }
-  if (opt === true || opt == null) return { ...ALL_ON }
+  if (opt === true || opt == null) return { ...LEVEL_PRESETS[2] }
   if (typeof opt === 'number') return { ...(LEVEL_PRESETS[opt] || ALL_ON) }
   if (typeof opt === 'object') {
     const baseLevel = typeof opt.level === 'number' ? opt.level : 2
     const base = LEVEL_PRESETS[baseLevel] || ALL_ON
     const out = { ...base }
-    for (const n of PASS_NAMES) if (n in opt) out[n] = !!opt[n]
+    for (const n of PASS_NAMES) if (n in opt) out[n] = n === 'nestedSmallConstForUnroll' && opt[n] === 'auto' ? 'auto' : !!opt[n]
     return out
   }
   return { ...ALL_ON }
@@ -137,8 +149,10 @@ export function hoistPtrType(fn) {
 
     if (op === 'call' && node[1] === '$__ptr_type' && node.length === 3) {
       const arg = node[2]
-      if (Array.isArray(arg) && arg[0] === 'local.get' && typeof arg[1] === 'string') {
-        const x = arg[1]
+      // Post-i64 migration: arg is (i64.reinterpret_f64 (local.get X)). Peel both wrappers.
+      const inner = (Array.isArray(arg) && arg[0] === 'i64.reinterpret_f64' && arg.length === 2) ? arg[1] : arg
+      if (Array.isArray(inner) && inner[0] === 'local.get' && typeof inner[1] === 'string') {
+        const x = inner[1]
         let region = open.get(x)
         if (!region) {
           region = []
@@ -449,9 +463,11 @@ export function hoistInvariantPtrOffset(fn) {
       const callee = node[1]
       if (callee === '$__ptr_offset' && node.length === 3) {
         const a = node[2]
-        if (Array.isArray(a) && a[0] === 'local.get' && typeof a[1] === 'string' && params.has(a[1])) {
-          let arr = sites.get(a[1])
-          if (!arr) { arr = []; sites.set(a[1], arr) }
+        // Post-i64 migration: arg may be (i64.reinterpret_f64 (local.get X)).
+        const inner = (Array.isArray(a) && a[0] === 'i64.reinterpret_f64' && a.length === 2) ? a[1] : a
+        if (Array.isArray(inner) && inner[0] === 'local.get' && typeof inner[1] === 'string' && params.has(inner[1])) {
+          let arr = sites.get(inner[1])
+          if (!arr) { arr = []; sites.set(inner[1], arr) }
           arr.push({ parent, idx: pi })
           return
         }
@@ -459,8 +475,9 @@ export function hoistInvariantPtrOffset(fn) {
       const isSafe = SAFE_OFFSET_CALLS.has(callee)
       for (let i = 2; i < node.length; i++) {
         const arg = node[i]
-        if (Array.isArray(arg) && arg[0] === 'local.get' && typeof arg[1] === 'string' && params.has(arg[1])) {
-          if (!isSafe) unsafe.add(arg[1])
+        const inner = (Array.isArray(arg) && arg[0] === 'i64.reinterpret_f64' && arg.length === 2) ? arg[1] : arg
+        if (Array.isArray(inner) && inner[0] === 'local.get' && typeof inner[1] === 'string' && params.has(inner[1])) {
+          if (!isSafe) unsafe.add(inner[1])
           continue
         }
         walk(arg, node, i)
@@ -497,7 +514,7 @@ export function hoistInvariantPtrOffset(fn) {
     if (arr.length < 2) continue
     const tLocal = `$__po${hoistId++}`
     newLocals.push(['local', tLocal, 'i32'])
-    snaps.push(['local.set', tLocal, ['call', '$__ptr_offset', ['local.get', X]]])
+    snaps.push(['local.set', tLocal, ['call', '$__ptr_offset', ['i64.reinterpret_f64', ['local.get', X]]]])
     for (const { parent, idx } of arr) {
       parent[idx] = ['local.get', tLocal]
     }
@@ -506,6 +523,132 @@ export function hoistInvariantPtrOffset(fn) {
   if (newLocals.length) fn.splice(bodyStart, 0, ...newLocals, ...snaps)
 }
 
+/**
+ * Per-loop hoist of `(call $__ptr_offset (local.get X))` for locals X that are
+ * loop-invariant (no `local.set/local.tee` to X anywhere in the loop body, no
+ * non-safe call inside the loop). Mirrors `hoistInvariantCellLoads`.
+ *
+ * The function-level pass above only handles params (single hoist at func
+ * entry). This pass handles assigned locals that go invariant within a loop
+ * scope — the aos pattern: `let rows = []; for (...) rows.push(...)` — outside
+ * the build loop, `rows` is no longer reassigned, so `__ptr_offset(rows)` is
+ * loop-invariant in the consumer loop.
+ *
+ * Inside-out: inner loops first. After an inner-loop hoist, the outer loop now
+ * contains a `(local.set $__pol (call $__ptr_offset ...))` whose call gets
+ * hoisted again at the outer level, climbing the snap up to the outermost loop
+ * where X is invariant. Cleanup of the chained `local.set` movs is handled by
+ * watr CSE/DCE.
+ */
+export function hoistInvariantPtrOffsetLoop(fn) {
+  if (!Array.isArray(fn) || fn[0] !== 'func') return
+  const bodyStart = findBodyStart(fn)
+  if (bodyStart < 0) return
+
+  let snapId = 0
+  while (fn.some(n => Array.isArray(n) && n[0] === 'local' && n[1] === `$__pol${snapId}`)) snapId++
+  const newLocals = []
+
+  // Refcount across the whole fn to skip shared subtrees (watr CSE may leave them).
+  const refcount = new Map()
+  const countRefs = (node) => {
+    if (!Array.isArray(node)) return
+    const n = (refcount.get(node) || 0) + 1
+    refcount.set(node, n)
+    if (n > 1) return
+    for (let i = 0; i < node.length; i++) countRefs(node[i])
+  }
+  countRefs(fn)
+
+  const processLoop = (loopNode) => {
+    // Recurse into inner loops first (bottom-up).
+    for (let i = 1; i < loopNode.length; i++) {
+      const child = loopNode[i]
+      if (!Array.isArray(child)) continue
+      processNode(child, loopNode, i)
+    }
+
+    // Scan loop body (including nested loops) for writes to any local and for
+    // any non-safe call. Non-safe calls could realloc/move the underlying
+    // array storage and invalidate cached offsets.
+    const writes = new Set()
+    let unsafe = false
+    const scan = (node) => {
+      if (!Array.isArray(node)) return
+      const op = node[0]
+      if (op === 'local.set' || op === 'local.tee') {
+        if (typeof node[1] === 'string') writes.add(node[1])
+        for (let i = 2; i < node.length; i++) scan(node[i])
+        return
+      }
+      if (op === 'call') {
+        if (!SAFE_OFFSET_CALLS.has(node[1])) unsafe = true
+        for (let i = 2; i < node.length; i++) scan(node[i])
+        return
+      }
+      if (op === 'call_ref' || op === 'call_indirect') {
+        unsafe = true
+        for (let i = 1; i < node.length; i++) scan(node[i])
+        return
+      }
+      for (let i = 1; i < node.length; i++) scan(node[i])
+    }
+    for (let i = 1; i < loopNode.length; i++) scan(loopNode[i])
+    if (unsafe) return []
+
+    // Collect call sites for invariant locals. Skip nested loops (already
+    // processed) but recurse through everything else.
+    const sites = new Map()  // localName → [{ parent, idx }]
+    const collect = (node, parent, idx) => {
+      if (!Array.isArray(node)) return
+      const op = node[0]
+      if (op === 'loop') return
+      if (op === 'call' && node[1] === '$__ptr_offset' && node.length === 3) {
+        const a = node[2]
+        const inner = (Array.isArray(a) && a[0] === 'i64.reinterpret_f64' && a.length === 2) ? a[1] : a
+        if (Array.isArray(inner) && inner[0] === 'local.get' && typeof inner[1] === 'string'
+            && !writes.has(inner[1])
+            && (refcount.get(node) || 0) <= 1
+            && (refcount.get(parent) || 0) <= 1) {
+          let arr = sites.get(inner[1])
+          if (!arr) { arr = []; sites.set(inner[1], arr) }
+          arr.push({ parent, idx })
+        }
+        return
+      }
+      for (let i = 0; i < node.length; i++) collect(node[i], node, i)
+    }
+    for (let i = 1; i < loopNode.length; i++) collect(loopNode[i], loopNode, i)
+
+    const snaps = []
+    for (const [X, arr] of sites) {
+      if (arr.length < 1) continue
+      const snapName = `$__pol${snapId++}`
+      newLocals.push(['local', snapName, 'i32'])
+      snaps.push(['local.set', snapName, ['call', '$__ptr_offset', ['i64.reinterpret_f64', ['local.get', X]]]])
+      for (const { parent, idx } of arr) {
+        parent[idx] = ['local.get', snapName]
+      }
+    }
+    return snaps
+  }
+
+  const processNode = (node, parent, idx) => {
+    if (!Array.isArray(node)) return
+    const op = node[0]
+    if (op === 'loop') {
+      const snaps = processLoop(node)
+      if (snaps.length) parent.splice(idx, 0, ...snaps)
+      return
+    }
+    for (let i = 0; i < node.length; i++) processNode(node[i], node, i)
+  }
+
+  for (let i = bodyStart; i < fn.length; i++) processNode(fn[i], fn, i)
+
+  if (newLocals.length) fn.splice(bodyStart, 0, ...newLocals)
+}
+
 /**
  * Hoist loop-invariant boxed-cell reads out of loops.
  *
@@ -665,6 +808,314 @@ export function hoistInvariantCellLoads(fn) {
   if (newLocals.length) fn.splice(bodyStart, 0, ...newLocals)
 }
 
+/**
+ * CSE for `(f64.load offset=K (local.get $X))` over straight-line regions
+ * where $X is an i32-typed local (an unboxed pointer in jz's value model).
+ *
+ * Aos hot path: `let p = rows[i]; xs[i] = p.x + p.y*0.25 + r;
+ *                ys[i] = p.y - p.z*0.5;
+ *                zs[i] = p.z + p.x*0.125`
+ * — emits 6 f64.load on $p (each of x/y/z twice); collapses to 3 unique loads
+ * shared via tee'd snap locals.
+ *
+ * Safety: jz's invariant — distinct unboxed-pointer locals come from distinct
+ * fresh allocations (analyzePtrUnboxable refuses to unbox aliased locals).
+ * So `(f64.store ADDR ...)` with base `(local.get $Y)` for $Y ≠ $X cannot
+ * touch addresses reachable via `$X + K`. Stores to typed-array slots in the
+ * loop body don't invalidate row-pointer reads.
+ *
+ * Region boundaries that flush the table:
+ *   - branch (br/br_if/br_table/return/unreachable)
+ *   - non-pure call
+ *   - loop / if  (control flow)
+ *   - local.set/local.tee on a tracked $X (invalidates that X's entries)
+ *   - store whose address tree references a tracked $X
+ * Blocks are treated as transparent — recurse into children.
+ */
+export function cseScalarLoad(fn) {
+  // DISABLED: the safety claim above relies on `analyzePtrUnboxable` having vetted
+  // every i32 local as a non-aliased fresh-allocation pointer. But this pass scans
+  // *all* i32 locals from `(local … i32)` decls — wasm-native i32 scalars (lengths,
+  // indices), narrow-ABI helper returns, and analyze.js's new arrayElemSchema-driven
+  // unboxes share the same declaration form. The metacircular path (jz-compiled
+  // watr.wasm) trips this: CSE'd loads survive across stores that legitimately
+  // mutate the same memory through a different i32 local, returning stale bytes
+  // (manifests as "memory access out of bounds" once a corrupted offset is
+  // dereferenced). Re-enable once candidacy is restricted to vetted pointers
+  // (e.g. emit-side annotation or rep-derived whitelist).
+  return
+  if (!Array.isArray(fn) || fn[0] !== 'func') return
+  const bodyStart = findBodyStart(fn)
+  if (bodyStart < 0) return
+
+  const i32Locals = new Set()
+  for (let i = 2; i < fn.length; i++) {
+    const c = fn[i]
+    if (Array.isArray(c) && (c[0] === 'local' || c[0] === 'param') && typeof c[1] === 'string' && c[2] === 'i32') {
+      i32Locals.add(c[1])
+    }
+  }
+  if (!i32Locals.size) return
+
+  let snapId = 0
+  while (fn.some(n => Array.isArray(n) && n[0] === 'local' && n[1] === `$__cs${snapId}`)) snapId++
+  const newLocals = []
+
+  // CSE table: key `${X}|${K}` → { snapName | null, anchorParent, anchorIdx }
+  const table = new Map()
+
+  const invalidateLocal = (X) => {
+    for (const key of table.keys()) {
+      if (key.startsWith(`${X}|`)) table.delete(key)
+    }
+  }
+
+  // Scan a node's subtree and return the set of i32 locals referenced via local.get.
+  const collectGets = (node, out) => {
+    if (!Array.isArray(node)) return
+    if (node[0] === 'local.get' && typeof node[1] === 'string' && i32Locals.has(node[1])) {
+      out.add(node[1])
+      return
+    }
+    for (let i = 1; i < node.length; i++) collectGets(node[i], out)
+  }
+
+  // Parse f64.load shape; returns { K, addrIdx } or null.
+  const parseLoad = (node) => {
+    if (!Array.isArray(node) || node[0] !== 'f64.load') return null
+    let K = 0, addrIdx = 1
+    if (typeof node[1] === 'string' && node[1].startsWith('offset=')) {
+      K = parseInt(node[1].slice(7), 10) | 0
+      addrIdx = 2
+    }
+    if (node.length <= addrIdx) return null
+    return { K, addrIdx }
+  }
+
+  const walk = (node, parent, idx) => {
+    if (!Array.isArray(node)) return
+    const op = node[0]
+
+    // Control-flow boundaries: clear table.
+    if (op === 'br' || op === 'br_if' || op === 'br_table' || op === 'return' || op === 'unreachable') {
+      // Process args first (a br_if value, br arg, etc. could still benefit from current table)
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      table.clear()
+      return
+    }
+
+    if (op === 'loop' || op === 'if') {
+      // Save table state isn't useful; recurse with cleared table, then clear after.
+      const saved = new Map(table)
+      table.clear()
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      // After leaving compound, conservatively assume invalidation.
+      table.clear()
+      // Restore? No — restoring would be unsafe since the compound may have written.
+      saved.clear()
+      return
+    }
+
+    if (op === 'call') {
+      const callee = node[1]
+      // Process args first.
+      for (let i = 2; i < node.length; i++) walk(node[i], node, i)
+      if (!SAFE_OFFSET_CALLS.has(callee)) table.clear()
+      return
+    }
+
+    if (op === 'call_ref' || op === 'call_indirect') {
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      table.clear()
+      return
+    }
+
+    if (op === 'local.set' || op === 'local.tee') {
+      // Process value first.
+      for (let i = 2; i < node.length; i++) walk(node[i], node, i)
+      const X = node[1]
+      if (typeof X === 'string') invalidateLocal(X)
+      return
+    }
+
+    // Stores: process operands first; if address tree references any tracked X,
+    // invalidate that X's entries.
+    if (op === 'f64.store' || op === 'i32.store' || op === 'i64.store'
+        || op === 'i32.store8' || op === 'i32.store16'
+        || op === 'i64.store8' || op === 'i64.store16' || op === 'i64.store32'
+        || op === 'f32.store') {
+      // Address may be node[1] (raw) or node[2] (when node[1] is offset=/align= attr).
+      let addrIdx = 1
+      if (typeof node[1] === 'string' && (node[1].startsWith('offset=') || node[1].startsWith('align='))) {
+        addrIdx = 2
+      }
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      const dirty = new Set()
+      collectGets(node[addrIdx], dirty)
+      for (const X of dirty) invalidateLocal(X)
+      return
+    }
+
+    // f64.load: try CSE.
+    const lp = parseLoad(node)
+    if (lp) {
+      const addr = node[lp.addrIdx]
+      if (Array.isArray(addr) && addr[0] === 'local.get' && typeof addr[1] === 'string' && i32Locals.has(addr[1])) {
+        const X = addr[1]
+        const key = `${X}|${lp.K}`
+        const entry = table.get(key)
+        if (entry) {
+          if (!entry.snapName) {
+            const snapName = `$__cs${snapId++}`
+            entry.snapName = snapName
+            newLocals.push(['local', snapName, 'f64'])
+            // Wrap anchor with (local.tee $snap originalLoad).
+            const orig = entry.anchorParent[entry.anchorIdx]
+            entry.anchorParent[entry.anchorIdx] = ['local.tee', snapName, orig]
+          }
+          parent[idx] = ['local.get', entry.snapName]
+          return
+        } else {
+          table.set(key, { snapName: null, anchorParent: parent, anchorIdx: idx })
+          // Don't recurse; (local.get $X) has no children of interest.
+          return
+        }
+      }
+      // Non-CSE'able address; recurse to find inner loads.
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      return
+    }
+
+    // Default: recurse.
+    for (let i = 0; i < node.length; i++) walk(node[i], node, i)
+  }
+
+  for (let i = bodyStart; i < fn.length; i++) walk(fn[i], fn, i)
+
+  if (newLocals.length) fn.splice(bodyStart, 0, ...newLocals)
+}
+
+/**
+ * CSE for pure f64 binary ops on local-only operands.
+ *
+ * Mandelbrot loop: condition computes `(f64.mul $zx $zx)` and `(f64.mul $zy $zy)`;
+ * body recomputes both inside `tx = zx*zx - zy*zy + cx`. Pure ops on locals can't
+ * alias memory — only `local.set/tee X` invalidates entries referencing X. Unlike
+ * `cseScalarLoad`, br_if doesn't need to clear (no memory aliasing concern).
+ *
+ * Targets nodes of shape `(OP A B)` where OP ∈ {f64.mul, f64.add, f64.sub} and
+ * A,B ∈ `(local.get X)` | `(f64.const N)`. Commutative ops (mul, add) sort
+ * operand keys for canonical form.
+ *
+ * Region boundaries:
+ *   - `local.set/tee X` → invalidates entries referencing X
+ *   - `loop`, `if` → recurse with cleared table; clear after (compound may have written)
+ *   - `call`, `call_ref`, `call_indirect` → no clear (calls don't write locals directly;
+ *     the surrounding `local.set/tee` handles that)
+ *   - `br/br_if/br_table/return/unreachable` → NO clear (pure values still valid)
+ */
+export function csePureExpr(fn) {
+  if (!Array.isArray(fn) || fn[0] !== 'func') return
+  const bodyStart = findBodyStart(fn)
+  if (bodyStart < 0) return
+
+  let snapId = 0
+  while (fn.some(n => Array.isArray(n) && n[0] === 'local' && n[1] === `$__pe${snapId}`)) snapId++
+  const newLocals = []
+
+  const COMMUTATIVE = new Set(['f64.mul', 'f64.add'])
+  const TARGET_OPS = new Set(['f64.mul', 'f64.add', 'f64.sub'])
+
+  // Encode a leaf operand to a stable string key. Returns null if not pure-leaf.
+  const leafKey = (n) => {
+    if (!Array.isArray(n)) return null
+    if (n[0] === 'local.get' && typeof n[1] === 'string') return `L:${n[1]}`
+    if (n[0] === 'f64.const') return `C:${n[1]}`
+    return null
+  }
+
+  // table: key → { snapName | null, anchorParent, anchorIdx, locals: Set<string> }
+  const table = new Map()
+
+  const invalidateLocal = (X) => {
+    for (const [key, entry] of table) {
+      if (entry.locals.has(X)) table.delete(key)
+    }
+  }
+
+  const walk = (node, parent, idx) => {
+    if (!Array.isArray(node)) return
+    const op = node[0]
+
+    if (op === 'loop' || op === 'if') {
+      const saved = new Map(table)
+      table.clear()
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      table.clear()
+      saved.clear()
+      return
+    }
+
+    // `then`/`else` branches of an `if` are mutually exclusive at runtime —
+    // a snap tee cached in the `then` branch is unset when the `else` runs.
+    // Isolate per-branch tables so a sibling branch can't reach into another's
+    // CSE entries.
+    if (op === 'then' || op === 'else') {
+      table.clear()
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      table.clear()
+      return
+    }
+
+    if (op === 'call' || op === 'call_ref' || op === 'call_indirect') {
+      // Calls don't write locals; recurse, no clear.
+      for (let i = 1; i < node.length; i++) walk(node[i], node, i)
+      return
+    }
+
+    if (op === 'local.set' || op === 'local.tee') {
+      for (let i = 2; i < node.length; i++) walk(node[i], node, i)
+      const X = node[1]
+      if (typeof X === 'string') invalidateLocal(X)
+      return
+    }
+
+    // Try CSE on (OP A B) where A,B are pure leaves.
+    if (TARGET_OPS.has(op) && node.length === 3) {
+      const ka = leafKey(node[1])
+      const kb = leafKey(node[2])
+      if (ka && kb) {
+        const key = COMMUTATIVE.has(op) && ka > kb ? `${op}|${kb}|${ka}` : `${op}|${ka}|${kb}`
+        const entry = table.get(key)
+        if (entry) {
+          if (!entry.snapName) {
+            const snapName = `$__pe${snapId++}`
+            entry.snapName = snapName
+            newLocals.push(['local', snapName, 'f64'])
+            const orig = entry.anchorParent[entry.anchorIdx]
+            entry.anchorParent[entry.anchorIdx] = ['local.tee', snapName, orig]
+          }
+          parent[idx] = ['local.get', entry.snapName]
+          return
+        } else {
+          const locals = new Set()
+          if (ka.startsWith('L:')) locals.add(ka.slice(2))
+          if (kb.startsWith('L:')) locals.add(kb.slice(2))
+          table.set(key, { snapName: null, anchorParent: parent, anchorIdx: idx, locals })
+          return
+        }
+      }
+      // Fall through to recurse.
+    }
+
+    for (let i = 0; i < node.length; i++) walk(node[i], node, i)
+  }
+
+  for (let i = bodyStart; i < fn.length; i++) walk(fn[i], fn, i)
+
+  if (newLocals.length) fn.splice(bodyStart, 0, ...newLocals)
+}
+
 /**
  * Find the index of the first body-content child in a func node.
  * Skips `$name`, (export …), (import …), (type …), (param …), (result …), (local …).
@@ -698,7 +1149,10 @@ export function hoistConstantPool(funcs, addGlobal) {
     for (let i = 0; i < node.length; i++) {
       const c = node[i]
       if (Array.isArray(c) && c[0] === 'f64.const' && (typeof c[1] === 'number' || typeof c[1] === 'string')) {
-        const k = typeof c[1] === 'number' ? `n:${c[1]}` : `s:${c[1]}`
+        // Distinguish -0 from +0 by sign: template literal collapses both to "0".
+        const k = typeof c[1] === 'number'
+          ? (Object.is(c[1], -0) ? 'n:-0' : `n:${c[1]}`)
+          : `s:${c[1]}`
         counts.set(k, (counts.get(k) || 0) + 1)
         sites.push({ parent: node, idx: i, key: k })
       }
@@ -758,8 +1212,8 @@ export function specializeMkptr(funcs, addFunc, parseWat) {
   const SPECS = {
     '$__mkptr':     { params: ['i32', 'i32', 'i32'], result: 'f64', inline: true },
     '$__alloc_hdr': { params: ['i32', 'i32', 'i32'], result: 'i32' },
-    '$__typed_idx': { params: ['f64', 'i32'],        result: 'f64' },
-    '$__str_idx':   { params: ['f64', 'i32'],        result: 'f64' },
+    '$__typed_idx': { params: ['i64', 'i32'],        result: 'f64' },
+    '$__str_idx':   { params: ['i64', 'i32'],        result: 'f64' },
   }
   const MIN_USES = 5
 
@@ -815,9 +1269,9 @@ export function specializeMkptr(funcs, addFunc, parseWat) {
     // $__mkptr inline fast path: bake (type, aux) literals into i64.const template.
     if (target === '$__mkptr' && spec.inline && parts[0].startsWith('L:') && parts[1].startsWith('L:')) {
       const type = +parts[0].slice(2), aux = +parts[1].slice(2)
-      const tmpl = (NAN_PREFIX_BITS << 48n)
-        | ((BigInt(type) & 0xFn) << 47n)
-        | ((BigInt(aux) & 0x7FFFn) << 32n)
+      const tmpl = LAYOUT.NAN_PREFIX_BITS
+        | ((BigInt(type) & BigInt(LAYOUT.TAG_MASK)) << BigInt(LAYOUT.TAG_SHIFT))
+        | ((BigInt(aux) & BigInt(LAYOUT.AUX_MASK)) << BigInt(LAYOUT.AUX_SHIFT))
       // Third arg (offset) may also be literal — emit (f64.const nan:…) then.
       if (parts[2].startsWith('L:')) {
         // Fully literal: all sites can be f64.const — no helper needed, handled in rewrite below.
@@ -859,10 +1313,10 @@ export function specializeMkptr(funcs, addFunc, parseWat) {
     // $__mkptr fully literal (rare — mkPtrIR usually folds these ahead of us, but defensive):
     if (target === '$__mkptr' && parts[0].startsWith('L:') && parts[1].startsWith('L:') && parts[2].startsWith('L:')) {
       const type = +parts[0].slice(2), aux = +parts[1].slice(2), off = +parts[2].slice(2)
-      const bits = (NAN_PREFIX_BITS << 48n)
-        | ((BigInt(type) & 0xFn) << 47n)
-        | ((BigInt(aux) & 0x7FFFn) << 32n)
-        | (BigInt(off >>> 0) & 0xFFFFFFFFn)
+      const bits = LAYOUT.NAN_PREFIX_BITS
+        | ((BigInt(type) & BigInt(LAYOUT.TAG_MASK)) << BigInt(LAYOUT.TAG_SHIFT))
+        | ((BigInt(aux) & BigInt(LAYOUT.AUX_MASK)) << BigInt(LAYOUT.AUX_SHIFT))
+        | (BigInt(off >>> 0) & BigInt(LAYOUT.OFFSET_MASK))
       const n = ['f64.const', 'nan:0x' + bits.toString(16).toUpperCase().padStart(16, '0')]
       n.type = 'f64'
       parent[idx] = n
@@ -1053,16 +1507,28 @@ export function sortStrPoolByFreq(funcs, strPoolRef, strDedupMap) {
 export function optimizeFunc(fn, cfg) {
   if (cfg && cfg.hoistPtrType === false &&
       cfg.hoistInvariantPtrOffset === false &&
+      cfg.hoistInvariantPtrOffsetLoop === false &&
       cfg.fusedRewrite === false &&
       cfg.hoistAddrBase === false &&
       cfg.hoistInvariantCellLoads === false &&
-      cfg.sortLocalsByUse === false) return
+      cfg.cseScalarLoad === false &&
+      cfg.csePureExpr === false &&
+      cfg.sortLocalsByUse === false &&
+      cfg.vectorizeLaneLocal === false) return
   if (!cfg || cfg.hoistPtrType !== false) hoistPtrType(fn)
   if (!cfg || cfg.hoistInvariantPtrOffset !== false) hoistInvariantPtrOffset(fn)
+  if (!cfg || cfg.hoistInvariantPtrOffsetLoop !== false) hoistInvariantPtrOffsetLoop(fn)
   const counts = new Map()
   if (!cfg || cfg.fusedRewrite !== false) fusedRewrite(fn, counts)
   if (!cfg || cfg.hoistAddrBase !== false) hoistAddrBase(fn)
   if (!cfg || cfg.hoistInvariantCellLoads !== false) hoistInvariantCellLoads(fn)
+  if (!cfg || cfg.cseScalarLoad !== false) cseScalarLoad(fn)
+  if (!cfg || cfg.csePureExpr !== false) csePureExpr(fn)
+  // Vectorizer runs only in the POST-watr phase (`__phase === 'post'`), where
+  // narrow.js + watr have produced the canonical lane-local shape. Running it
+  // pre-watr matches a noisier IR and would re-fire post-watr on the residual
+  // tail — producing nested SIMD blocks. Single-pass via phase gating.
+  if (cfg && cfg.vectorizeLaneLocal === true && cfg.__phase === 'post') vectorizeLaneLocal(fn)
   if (!cfg || cfg.sortLocalsByUse !== false) sortLocalsByUse(fn, cfg && cfg.fusedRewrite !== false ? counts : null)
 }
 
@@ -1109,32 +1575,59 @@ function walkRewrite(node, doInline, counts) {
   if (doInline && op === 'call' && node.length === 3 && typeof node[1] === 'string') {
     const fname = node[1]
     if (fname === '$__ptr_type') return ['i32.and',
-      ['i32.wrap_i64', ['i64.shr_u', ['i64.reinterpret_f64', node[2]], ['i64.const', 47]]],
-      ['i32.const', 0xF]]
+      ['i32.wrap_i64', ['i64.shr_u', node[2], ['i64.const', LAYOUT.TAG_SHIFT]]],
+      ['i32.const', LAYOUT.TAG_MASK]]
     if (fname === '$__ptr_aux') return ['i32.and',
-      ['i32.wrap_i64', ['i64.shr_u', ['i64.reinterpret_f64', node[2]], ['i64.const', 32]]],
-      ['i32.const', 0x7FFF]]
-    if (fname === '$__is_null') return ['i64.eq', ['i64.reinterpret_f64', node[2]], ['i64.const', '0x7FF8000100000000']]
-    if (fname === '$__is_nullish' && Array.isArray(node[2]) && node[2][0] === 'local.get') return ['i32.or',
-      ['i64.eq', ['i64.reinterpret_f64', node[2]], ['i64.const', '0x7FF8000100000000']],
-      ['i64.eq', ['i64.reinterpret_f64', node[2]], ['i64.const', '0x7FF8000000000001']]]
-    if (fname === '$__is_truthy' && Array.isArray(node[2]) && node[2][0] === 'local.get') {
-      const lget = node[2]
-      const bits = ['i64.reinterpret_f64', lget]
+      ['i32.wrap_i64', ['i64.shr_u', node[2], ['i64.const', LAYOUT.AUX_SHIFT]]],
+      ['i32.const', LAYOUT.AUX_MASK]]
+    if (fname === '$__is_null') return ['i64.eq', node[2], ['i64.const', NULL_BITS]]
+    if (fname === '$__is_nullish' && Array.isArray(node[2]) && node[2][0] === 'i64.reinterpret_f64'
+        && Array.isArray(node[2][1]) && node[2][1][0] === 'local.get') return ['i32.or',
+      ['i64.eq', node[2], ['i64.const', NULL_BITS]],
+      ['i64.eq', node[2], ['i64.const', UNDEF_BITS]]]
+    if (fname === '$__is_truthy' && Array.isArray(node[2]) && node[2][0] === 'i64.reinterpret_f64'
+        && Array.isArray(node[2][1]) && node[2][1][0] === 'local.get') {
+      const lget = node[2][1]
+      const bits = node[2]
       return ['if', ['result', 'i32'],
         ['f64.eq', lget, lget],
         ['then', ['f64.ne', lget, ['f64.const', 0]]],
         ['else', ['i32.and',
           ['i32.and',
-            ['i64.ne', bits, ['i64.const', '0x7FF8000000000000']],
-            ['i64.ne', bits, ['i64.const', '0x7FF8000100000000']]],
+            ['i64.ne', bits, ['i64.const', NAN_BITS]],
+            ['i64.ne', bits, ['i64.const', NULL_BITS]]],
           ['i32.and',
-            ['i64.ne', bits, ['i64.const', '0x7FF8000000000001']],
-            ['i64.ne', bits, ['i64.const', '0x7FFA800000000000']]]]]]
+            ['i64.ne', bits, ['i64.const', UNDEF_BITS]],
+            ['i64.ne', bits, ['i64.const', '0x7FFA400000000000']]]]]]
     }
   }
 
   // Peephole: rebox/unbox round-trips
+  if ((op === 'f64.convert_i32_s' || op === 'f64.convert_i32_u') && node.length === 2) {
+    const a = node[1]
+    if (Array.isArray(a) && a[0] === 'i32.const') {
+      const n = typeof a[1] === 'number' ? a[1] : typeof a[1] === 'string' ? Number(a[1]) : NaN
+      if (Number.isFinite(n)) return ['f64.const', op === 'f64.convert_i32_u' ? n >>> 0 : n]
+    }
+  }
+  if (op === 'f64.mul' && node.length === 3) {
+    const a = node[1], b = node[2]
+    const isTwo = x => Array.isArray(x) && x[0] === 'f64.const' && x[1] === 2
+    const isCheapF64 = x => Array.isArray(x) &&
+      ((x[0] === 'local.get' && typeof x[1] === 'string') ||
+       (x[0] === 'f64.const' && typeof x[1] === 'number'))
+    if (isTwo(a) && isCheapF64(b)) return ['f64.add', b, b]
+    if (isTwo(b) && isCheapF64(a)) return ['f64.add', a, a]
+  }
+  if (op === 'i32.trunc_sat_f64_s' && node.length === 2) {
+    const a = node[1]
+    if (Array.isArray(a) && a[0] === 'f64.convert_i32_s' && a.length === 2) return a[1]
+  }
+  if (op === 'i64.trunc_sat_f64_s' && node.length === 2) {
+    const a = node[1]
+    if (Array.isArray(a) && a[0] === 'f64.convert_i32_s' && a.length === 2) return ['i64.extend_i32_s', a[1]]
+    if (Array.isArray(a) && a[0] === 'f64.convert_i32_u' && a.length === 2) return ['i64.extend_i32_u', a[1]]
+  }
   if (op === 'i64.reinterpret_f64' && node.length === 2) {
     const a = node[1]
     if (Array.isArray(a) && a[0] === 'f64.reinterpret_i64' && a.length === 2) return a[1]
@@ -1350,3 +1843,96 @@ export function sortLocalsByUse(fn, precomputedCounts) {
   locals.sort((a, b) => (counts.get(b[1]) || 0) - (counts.get(a[1]) || 0))
   localIdxs.forEach((i, k) => { fn[i] = locals[k] })
 }
+
+/**
+ * Module-level arena rewind: transitive escape analysis.
+ *
+ * Per-function `applyArenaRewind` in compile.js is limited to a static whitelist
+ * of internal helpers. This pass generalizes by building a call graph and
+ * propagating "arena-safe callee" status via fixed-point iteration.
+ *
+ * A function is an arena-safe callee if:
+ *   - no global.set, call_indirect, call_ref in body
+ *   - all user-function calls are to other arena-safe callees
+ *
+ * A function is arena-rewindable (gets heap save/restore injected) if:
+ *   - single scalar result (f64 or i32)
+ *   - contains allocation ($__alloc / $__alloc_hdr)
+ *   - no global.set, return_call, call_indirect, call_ref
+ *   - all user-function calls are to arena-safe callees
+ *   - does NOT return a pointer (checked via ptrTypes map from compile.js)
+ *
+ * Unlike the per-function pass, this does NOT require 0 params.
+ *
+ * @param {Array[]} fns - Array of func IR nodes (sec.funcs + sec.stdlib + sec.start)
+ * @param {boolean} sharedMemory - Whether memory is shared (affects heap get/set IR)
+ * @param {Map<string, {ptrKind: *}|null>} [ptrTypes] - Map from func name to ptrKind info.
+ *   Functions with ptrKind != null return pointers and cannot be rewound.
+ *   If omitted, no pointer-return check is done (conservative: fewer functions rewound).
+ */
+export function arenaRewindModule(fns) {
+  const BUILTIN_SAFE = new Set([
+    '$__alloc', '$__alloc_hdr', '$__mkptr',
+    '$__ptr_offset', '$__ptr_type', '$__ptr_aux',
+    '$__len', '$__cap', '$__typed_shift', '$__typed_data',
+  ])
+
+  // Phase 1: collect per-function metadata
+  const fnMap = new Map()
+  for (const fn of fns) {
+    if (!Array.isArray(fn) || fn[0] !== 'func') continue
+    const name = fn[1]
+    if (typeof name !== 'string') continue
+
+    let results = [], hasGlobalSet = false, hasReturnCall = false
+    let hasCallIndirect = false, hasCallRef = false, hasAlloc = false
+    const calls = new Set()
+    const bodyStart = findBodyStart(fn)
+
+    for (let i = 2; i < fn.length; i++) {
+      const c = fn[i]
+      if (!Array.isArray(c)) continue
+      if (c[0] === 'result') { results.push(c[1] || c[2]); continue }
+      if (i >= bodyStart) break
+    }
+
+    const scan = node => {
+      if (!Array.isArray(node)) return
+      const op = node[0]
+      if (op === 'global.set') hasGlobalSet = true
+      else if (op === 'return_call') hasReturnCall = true
+      else if (op === 'call_indirect') hasCallIndirect = true
+      else if (op === 'call_ref') hasCallRef = true
+      else if (op === 'call') {
+        const callee = node[1]
+        if (callee === '$__alloc' || callee === '$__alloc_hdr') hasAlloc = true
+        if (typeof callee === 'string' && !BUILTIN_SAFE.has(callee)) calls.add(callee)
+      }
+      for (let i = 1; i < node.length; i++) scan(node[i])
+    }
+    for (let i = bodyStart; i < fn.length; i++) scan(fn[i])
+
+    fnMap.set(name, {
+      fn, results,
+      hasGlobalSet, hasReturnCall, hasCallIndirect, hasCallRef, hasAlloc,
+      calls: [...calls],
+    })
+  }
+
+  // Phase 2: fixed-point transitive safety analysis
+  const safeCallees = new Set(BUILTIN_SAFE)
+  let changed = true
+  while (changed) {
+    changed = false
+    for (const [name, info] of fnMap) {
+      if (safeCallees.has(name)) continue
+      if (info.hasGlobalSet || info.hasCallIndirect || info.hasCallRef) continue
+      if (info.calls.every(c => safeCallees.has(c) || !fnMap.has(c))) {
+        safeCallees.add(name)
+        changed = true
+      }
+    }
+  }
+
+  return safeCallees
+}