jz 0.5.1 → 0.6.0

package/src/{assemble.js → wat/assemble.js}

@@ -12,8 +12,8 @@
  * @module assemble
  */
 
-import { parse as parseWat } from 'watr'
-import { ctx, inc, resolveIncludes, err, PTR, LAYOUT } from './ctx.js'
+import parseWat from 'watr/parse'
+import { ctx, inc, resolveIncludes, err, PTR, LAYOUT, HEAP, declGlobal } from '../ctx.js'
 
 // Stdlib WAT templates are fixed text (or feature-keyed text from a factory) —
 // `parseWat` of the same string always yields the same tree. Parsing is the
@@ -33,10 +33,12 @@ const parseTemplate = (str) => {
   if (tmpl === undefined) stdlibParseCache.set(str, tmpl = parseWat(str))
   return cloneTemplate(tmpl)
 }
-import { T, VAL, analyzeValTypes } from './analyze.js'
-import { optimizeFunc, collectVolatileGlobals, hoistConstantPool, specializeMkptr, specializePtrBase, sortStrPoolByFreq, arenaRewindModule } from './optimize.js'
-import { emit } from './emit.js'
-import { mkPtrIR, MAX_CLOSURE_ARITY, MEM_OPS, findBodyStart } from './ir.js'
+import { T } from '../ast.js'
+import { analyzeValTypes, analyzeBody } from '../compile/analyze.js'
+import { VAL } from '../reps.js'
+import { optimizeFunc, collectVolatileGlobals, collectReachableGlobalWrites, hoistGlobalPtrOffset, stablePtrGlobalNames, hoistConstantPool, specializeMkptr, specializePtrBase, sortStrPoolByFreq, arenaRewindModule } from '../optimize/index.js'
+import { emit } from '../compile/emit.js'
+import { mkPtrIR, MAX_CLOSURE_ARITY, MEM_OPS, findBodyStart } from '../ir.js'
 
 // NaN-prefix top-13-bits as BigInt — used by the static-prefix-strip pass
 const NAN_PREFIX = BigInt(LAYOUT.NAN_PREFIX)
@@ -46,16 +48,16 @@ const TAG_SHIFT_BIG = BigInt(LAYOUT.TAG_SHIFT)
 const AUX_SHIFT_BIG = BigInt(LAYOUT.AUX_SHIFT)
 const SSO_BIT_BIG = BigInt(LAYOUT.SSO_BIT)
 
-// memory[1020] holds the heap pointer only for shared memory (wasm globals are
+// memory[HEAP.PTR_ADDR] holds the heap pointer only for shared memory (wasm globals are
 // per-instance — see module/core.js comment). Non-shared memory uses $__heap.
 const heapUsesMem = () => ctx.memory.shared
 
 const heapGetIR = () => heapUsesMem()
-  ? ['i32.load', ['i32.const', 1020]]
+  ? ['i32.load', ['i32.const', HEAP.PTR_ADDR]]
   : ['global.get', '$__heap']
 
 const heapSetIR = value => heapUsesMem()
-  ? ['i32.store', ['i32.const', 1020], value]
+  ? ['i32.store', ['i32.const', HEAP.PTR_ADDR], value]
   : ['global.set', '$__heap', value]
 
 const ARENA_SAFE_CALLS = new Set([
@@ -134,8 +136,20 @@ export function buildStartFn(ast, sec, closureFuncs, compilePendingClosures) {
   ctx.func.locals = new Map()
   ctx.func.localReps = null
   ctx.func.boxed = new Map()
+  ctx.func.cellTypes = new Set()
   ctx.func.stack = []
   ctx.func.current = { params: [], results: [] }
+  // Reserve prepare-generated temp names (for-of `arrVar`/`idx`/`len`,
+  // destructure scratch, …) in the __start frame so emit's temp()/tempI32()
+  // skip them — the same pre-seed analyzeFuncForEmit gives every function frame.
+  // Only T-sentinel names: they're always __start locals (user module-scope
+  // bindings become globals and can't contain T). Without this, prepare's
+  // `${T}arr${n}` collides with an emit-time tempI32('arr') at the same uniq,
+  // declaring the array pointer's local i32 and corrupting it via convert_i32_s.
+  const seedGeneratedLocals = (body) => {
+    for (const [n, t] of analyzeBody(body).locals)
+      if (n.includes(T) && !ctx.func.locals.has(n)) ctx.func.locals.set(n, t)
+  }
   analyzeValTypes(ast)
   const normalizeIR = ir => !ir?.length ? [] : Array.isArray(ir[0]) ? ir : [ir]
 
@@ -143,9 +157,11 @@ export function buildStartFn(ast, sec, closureFuncs, compilePendingClosures) {
   if (ctx.module.moduleInits) {
     for (const mi of ctx.module.moduleInits) {
       analyzeValTypes(mi)
+      seedGeneratedLocals(mi)
       moduleInits.push(...normalizeIR(emit(mi)))
     }
   }
+  seedGeneratedLocals(ast)
   const init = emit(ast)
 
   // Module-scope object literals can create closure bodies while `emit(ast)`
@@ -342,8 +358,12 @@ export function finalizeClosureTable(sec) {
   }
   for (const fn of sec.funcs) { scan(fn); if (indirectUsed) break }
   if (!indirectUsed) for (const fn of sec.start) scan(fn)
-  if (!indirectUsed) for (const s of Object.keys(ctx.core.stdlib)) {
-    if (ctx.core.stdlib[s]?.includes?.('call_indirect')) { indirectUsed = true; break }
+  // stdlib values are mixed: WAT-template strings + lazy generator functions.
+  // Only the string templates can carry a literal `call_indirect`; a typeof
+  // guard skips the generators (where `.includes` is meaningless — and on a jz
+  // closure receiver would read the closure pointer as a string, out of bounds).
+  if (!indirectUsed) for (const tpl of Object.values(ctx.core.stdlib)) {
+    if (typeof tpl === 'string' && tpl.includes('call_indirect')) { indirectUsed = true; break }
   }
   if (indirectUsed) {
     if (!ctx.closure.table) ctx.closure.table = []
@@ -409,7 +429,11 @@ export 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', '__clear']) if (!ctx.core.includes.has(fn)) ctx.core.includes.add(fn)
+    for (const fn of ['__alloc', '__alloc_hdr', '__clear']) ctx.core.includes.add(fn)
+    // Late-add of allocators may pull in transitive deps (__alloc → __memgrow,
+    // etc.) that the initial resolveIncludes did not yet see; re-resolve.
+    // No-op when the alloc trio was already present.
+    resolveIncludes()
     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])
@@ -443,31 +467,34 @@ export function syncImports(sec) {
 /**
  * Phase: whole-module + per-function optimization passes.
  */
-export function optimizeModule(sec) {
+export function optimizeModule(sec, profiler) {
+  const t = profiler?.time ? (name, fn) => profiler.time(`optMod:${name}`, fn) : (_, fn) => fn()
   const cfg = ctx.transform.optimize
-  if (!cfg || cfg.specializeMkptr !== false)
-    specializeMkptr([...sec.funcs, ...sec.stdlib, ...sec.start], wat => sec.stdlib.push(parseWat(wat)), parseWat)
-  if (!cfg || cfg.specializePtrBase !== false)
-    specializePtrBase([...sec.funcs, ...sec.stdlib, ...sec.start], wat => sec.stdlib.push(parseWat(wat)), parseWat)
-  if (ctx.runtime.strPool && (!cfg || cfg.sortStrPoolByFreq !== false)) {
+  if (!cfg || cfg.specializeMkptr !== false) t('specializeMkptr', () =>
+    specializeMkptr([...sec.funcs, ...sec.stdlib, ...sec.start], wat => sec.stdlib.push(parseWat(wat)), parseWat))
+  if (!cfg || cfg.specializePtrBase !== false) t('specializePtrBase', () =>
+    specializePtrBase([...sec.funcs, ...sec.stdlib, ...sec.start], wat => sec.stdlib.push(parseWat(wat)), parseWat))
+  if (ctx.runtime.strPool && (!cfg || cfg.sortStrPoolByFreq !== false)) t('sortStrPool', () => {
     const poolRef = { pool: ctx.runtime.strPool }
     sortStrPoolByFreq([...sec.funcs, ...sec.stdlib, ...sec.start], poolRef, ctx.runtime.strPoolDedup)
     ctx.runtime.strPool = poolRef.pool
-  }
-  // Backfill globalTypes for runtime globals declared only in ctx.scope.globals
-  // (e.g., __schema_tbl, __strBase). Parses the WAT string to infer i32/f64/i64.
-  if (ctx.scope.globals) {
-    for (const [name, wat] of ctx.scope.globals) {
-      if (!wat || ctx.scope.globalTypes.has(name)) continue
-      const m = wat.match(/\(global\s+\$?\S+\s+(?:\(mut\s+)?(i32|i64|f64|f32)/)
-      if (m) ctx.scope.globalTypes.set(name, m[1])
-    }
-  }
+  })
+  // (globalTypes backfill gone: declGlobal sets the type at declaration.)
   // Build global name→type map from ctx.scope.globalTypes (keys without $) for promoteGlobals
   const globalTypesMap = ctx.scope.globalTypes ? new Map([...ctx.scope.globalTypes].map(([k, v]) => [`$${k}`, v])) : null
   const allFuncs = [...sec.funcs, ...sec.stdlib, ...sec.start]
-  const volatileGlobals = collectVolatileGlobals(allFuncs)
-  for (const s of allFuncs) optimizeFunc(s, cfg, globalTypesMap, volatileGlobals)
+  const volatileGlobals = t('volatileGlobals', () => collectVolatileGlobals(allFuncs))
+  const reachableWrites = t('reachableWrites', () => collectReachableGlobalWrites(allFuncs))
+  // Offset-hoist BEFORE promoteGlobals (inside optimizeFunc): value-promoting a
+  // stable-pointee global to a $_pg local would destroy the global.get pattern
+  // this pass matches, reverting rfft/diffusion to per-iteration resolves. After
+  // the hoist, the surviving global.get count is 1 (the entry snap) — naturally
+  // below promoteGlobals' threshold, so the two passes compose either way.
+  if (!cfg || cfg.hoistGlobalPtrOffset !== false) t('hoistGlobalPtr', () => {
+    const stable = stablePtrGlobalNames()
+    if (stable.size) for (const s of allFuncs) hoistGlobalPtrOffset(s, stable, reachableWrites)
+  })
+  t('optimizeFuncs', () => { for (const s of allFuncs) optimizeFunc(s, cfg, globalTypesMap, volatileGlobals, 'pre', reachableWrites) })
   if (!cfg || cfg.arenaRewind !== false) {
     const safeCallees = arenaRewindModule([...sec.funcs, ...sec.stdlib, ...sec.start])
     const fnByName = new Map()
@@ -481,11 +508,7 @@ export function optimizeModule(sec) {
     }
   }
   if (!cfg || cfg.hoistConstantPool !== false)
-    hoistConstantPool([...sec.funcs, ...sec.stdlib, ...sec.start], (name, wat) => {
-      ctx.scope.globals.set(name, wat)
-      const m = wat.match(/\(global\s+\$?\S+\s+(?:\(mut\s+)?(i32|i64|f64|f32)/)
-      if (m) ctx.scope.globalTypes.set(name, m[1])
-    })
+    hoistConstantPool([...sec.funcs, ...sec.stdlib, ...sec.start], (name, lit) => declGlobal(name, 'f64', lit))
 
   // Second promoteGlobals pass disabled: promoting hoistConstantPool's __fc*
   // globals regressed the watr perf micro-pin (WASM compile time increased).
@@ -502,12 +525,8 @@ export function optimizeModule(sec) {
     const heapBase = (dataLen + 7) & ~7
     // Non-shared memory always carries a $__heap global — start it past the
     // static data so the bump allocator never overwrites a literal.
-    ctx.scope.globals.set('__heap', `(global $__heap (export "__heap") (mut i32) (i32.const ${heapBase}))`)
-    ctx.scope.globalTypes.set('__heap', 'i32')
-    if (ctx.scope.globals.has('__heap_start')) {
-      ctx.scope.globals.set('__heap_start', `(global $__heap_start (mut i32) (i32.const ${heapBase}))`)
-      ctx.scope.globalTypes.set('__heap_start', 'i32')
-    }
+    declGlobal('__heap', 'i32', heapBase, { export: '__heap' })
+    if (ctx.scope.globals.has('__heap_start')) declGlobal('__heap_start', 'i32', heapBase)
     for (const s of sec.stdlib)
       if (s[0] === 'func' && s[1] === '$__clear')
         for (let i = 2; i < s.length; i++)
@@ -528,18 +547,33 @@ export function stripStaticDataPrefix(sec) {
   for (let i = 0; i < data.length; i++) buf[i] = data.charCodeAt(i)
   const dv = new DataView(buf.buffer)
   if (ctx.runtime.staticPtrSlots) {
+    // u32-half reads/writes — DataView's BigInt accessors are unfaithful in the
+    // self-host kernel; the offset lives entirely in the LE low word and the
+    // tag/aux fields entirely in the high word, so plain number math suffices.
     for (const slotOff of ctx.runtime.staticPtrSlots) {
       if (slotOff < prefix) continue
-      const bits = dv.getBigUint64(slotOff, true)
-      if (((bits >> 48n) & 0xFFF8n) !== NAN_PREFIX) continue
-      const ty = Number((bits >> TAG_SHIFT_BIG) & TAG_MASK_BIG)
+      const hi = dv.getUint32(slotOff + 4, true)
+      if (((hi >>> 16) & 0xFFF8) !== LAYOUT.NAN_PREFIX) continue
+      const ty = (hi >>> 15) & 15
       if (!SHIFTABLE.has(ty)) continue
-      if (ty === PTR.STRING && ((bits >> AUX_SHIFT_BIG) & SSO_BIT_BIG)) continue
-      const off = Number(bits & OFFSET_MASK_BIG)
+      if (ty === PTR.STRING && ((hi >>> (LAYOUT.AUX_SHIFT - 32)) & LAYOUT.SSO_BIT)) continue
+      const off = dv.getUint32(slotOff, true)
       if (off < prefix) continue
-      const hi = bits & ~OFFSET_MASK_BIG
-      dv.setBigUint64(slotOff, hi | BigInt(off - prefix), true)
+      dv.setUint32(slotOff, off - prefix, true)
+    }
+  }
+  // The intern index (buildInternTable) stores raw static-string ptrs as u32
+  // slots — shift each occupied slot like every other static reference, and
+  // re-declare the (already-declared) base global at its post-strip position.
+  if (ctx.runtime.internTable) {
+    const { base, size } = ctx.runtime.internTable
+    for (let i = 0; i < size; i++) {
+      const slot = base + i * 8 + 4
+      const off = dv.getUint32(slot, true)
+      if (off >= prefix) dv.setUint32(slot, off - prefix, true)
     }
+    ctx.runtime.internTable.base = base - prefix
+    declGlobal('__internBase', 'i32', base - prefix, { mut: false })
   }
   let s = ''
   for (let i = prefix; i < buf.length; i++) s += String.fromCharCode(buf[i])
@@ -559,6 +593,10 @@ export function stripStaticDataPrefix(sec) {
           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 (typeof child[0] === 'string' && child[0].endsWith('.store') &&
+        Array.isArray(child[1]) && child[1][0] === 'i32.const' &&
+        typeof child[1][1] === 'number' && child[1][1] >= prefix) {
+        child[1][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