jz 0.6.0 → 0.8.0

package/src/compile/infer.js

@@ -46,7 +46,7 @@
 import { ctx } from '../ctx.js'
 import { collectParamNames, ASSIGN_OPS } from '../ast.js'
 import { analyzeValTypes, analyzeIntCertain } from './analyze.js'
-import { staticObjectProps } from '../static.js'
+import { staticObjectProps, staticArrayElems } from '../static.js'
 import { typedElemCtor } from '../type.js'
 import { ctorFromElemAux } from '../../layout.js'
 import { shapeOfObjectLiteralAst, valTypeOf } from '../kind.js'
@@ -341,6 +341,41 @@ export function recordGlobalRep(name, expr) {
   }
   const ctor = typedElemCtor(expr)
   if (ctor) (ctx.scope.globalTypedElem ||= new Map()).set(name, ctor)
+  // Module-level const array literal with a uniform element val-type (e.g. a numeric
+  // table `const FREQS = [261.63, …]`): record it so `FREQS[i]` reads in any using
+  // function are typed (NUMBER) rather than untyped. Without this an untyped element
+  // read makes `s += FREQS[i]` take the polymorphic +/ToString path — dragging the
+  // entire string runtime (~5 kB) into a kernel that uses no strings. A function-local
+  // array gets this from analyzeValTypes; a module-level one is invisible to the using
+  // function's body walk, so capture it here. Soundness for a later `FREQS[i]=…` is the
+  // read-site dynWriteVars guard in valTypeOf (kind.js) — this is just the literal fact.
+  if (vt === VAL.ARRAY) {
+    const elems = staticArrayElems(expr)
+    if (elems && elems.length && elems.every(e => e != null)) {
+      let common = valTypeOf(elems[0])
+      for (let k = 1; k < elems.length && common != null; k++)
+        if (valTypeOf(elems[k]) !== common) common = null
+      if (common != null) updateGlobalRep(name, { arrayElemValType: common })
+      // Array-of-arrays numeric table (`const C = [[0,4,7], …]`): also record the
+      // nested element kind so `C[i][j]` (and `ch = C[i]; ch[j]`) reads stay typed —
+      // the same string-runtime drop as the flat case, one level down. Single-level
+      // (mirrors analyzeValTypes' local arrElemElemValTypes); deeper nesting falls back.
+      if (common === VAL.ARRAY) {
+        let nested = null, seen = false, ok = true
+        for (const el of elems) {
+          const inner = staticArrayElems(el)
+          if (!inner || !inner.length || !inner.every(e => e != null)) { ok = false; break }
+          for (const ie of inner) {
+            const ivt = valTypeOf(ie)
+            if (!seen) { nested = ivt; seen = true }
+            else if (ivt !== nested) { ok = false; break }
+          }
+          if (!ok) break
+        }
+        if (ok && nested != null) updateGlobalRep(name, { arrayElemElemValType: nested })
+      }
+    }
+  }
   // Static-shape capture for module-level object literals — lets `{ ...G.path }`
   // resolve its source schema by walking the global rep's shape tree at the
   // spread site (see shape walk in analyze.js / resolveSchema in object.js).

package/src/compile/loop-divmod.js

@@ -0,0 +1,109 @@
+// Loop induction strength-reduction for `i % w` and `(i / w) | 0`.
+//
+// JS `%` and `/` are float ops; `i % w` with a RUNTIME divisor lowers to f64
+// (i % 0 is NaN, so it can't be soundly typed i32 — see type.js exprType `%`).
+// That makes the column `x = i % w` an f64 local, which cascades: every dependent
+// neighbour index becomes f64 and each typed-array access pays an f64→i32 convert.
+// Measured ~5× slower than V8, which speculates the divisor non-zero and uses i32.
+//
+// This pass replaces the per-iteration division with incremental i32 counters in a
+// unit-stride loop: the column `cx` increments by 1 and wraps to 0 at `w`, bumping
+// the row `cy`. No division survives in the body, so the whole index chain stays
+// i32. Fully sound — counters are pure increment (the divisor-zero question never
+// arises), and if `w == 0` the loop's `i < w*h` guard never admits an iteration, so
+// the one-time seed `(i%w)|0` (NaN→0) is unused.
+//
+// Recognized (post-prepare AST): a `while` whose body increments one IV `i` by +1
+// exactly once, reads `['%', i, w]` and/or `['|', ['/', i, w], LIT0]` with `w`
+// loop-invariant, and has no `continue` / closure-capture of i,w (which could
+// desync the counters). Number literals are sparse-array holes `[, v]` (n[0] is the
+// hole = undefined), so literal tests use `== null`; created literals are bare numbers.
+
+import { findMutations } from './analyze-scans.js'
+import { litN, unitIncVar, normalizeLoop, closureMutatedVars, rewriteBlocks, freshLoopId } from './loop-model.js'
+
+const isMod = (n, i, w) => Array.isArray(n) && n[0] === '%' && n[1] === i && n[2] === w
+const isFloorDiv = (n, i, w) =>
+  Array.isArray(n) && n[0] === '|' && litN(n[2], 0) &&
+  Array.isArray(n[1]) && n[1][0] === '/' && n[1][1] === i && n[1][2] === w
+
+const usesPattern = (n, i, w, pred) => Array.isArray(n) && (pred(n, i, w) || n.some(c => usesPattern(c, i, w, pred)))
+const replace = (n, i, w, cx, cy) =>
+  !Array.isArray(n) ? n : isMod(n, i, w) ? cx : isFloorDiv(n, i, w) ? cy : n.map(c => replace(c, i, w, cx, cy))
+// a `continue` that targets THIS loop (not one nested inside) — would skip the increment
+const hasOuterContinue = (n) => Array.isArray(n) &&
+  (n[0] === 'continue' || (n[0] !== 'while' && n[0] !== 'for' && n[0] !== 'do' && n[0] !== '=>' && n.some(hasOuterContinue)))
+
+
+// Try to strength-reduce one `while` statement. Returns [seed, loop] or null. `cm` is
+// the function's closure-mutated-vars set (an IV/divisor in it is unsafe).
+function tryReduce(stmt, cm) {
+  const L = normalizeLoop(stmt)
+  if (!L || L.kind !== 'while') return null
+  const cond = L.cond, lbody = L.body
+  if (!Array.isArray(lbody) || lbody[0] !== ';') return null
+
+  // exactly one IV incremented by +1
+  let iv = null, ivIdx = -1
+  for (let k = 1; k < lbody.length; k++) {
+    const v = unitIncVar(lbody[k])
+    if (v) { if (iv) return null; iv = v; ivIdx = k }
+  }
+  if (!iv) return null
+
+  // a single invariant divisor `w` (`i % w` / `(i/w)|0` all share it)
+  let w = null
+  const findW = (n) => {
+    if (!Array.isArray(n)) return
+    const d = n[0] === '%' && n[1] === iv ? n[2]
+      : (n[0] === '|' && litN(n[2], 0) && Array.isArray(n[1]) && n[1][0] === '/' && n[1][1] === iv) ? n[1][2] : null
+    if (typeof d === 'string') { if (w == null) w = d; else if (w !== d) w = false }
+    n.forEach(findW)
+  }
+  findW(lbody)
+  if (!w || w === iv) return null
+
+  const usesMod = usesPattern(lbody, iv, w, isMod)
+  const usesDiv = usesPattern(lbody, iv, w, isFloorDiv)
+  if (!usesMod && !usesDiv) return null
+
+  // soundness: w invariant, iv written ONLY by the increment, no continue/closure capture
+  const wMut = new Set(); findMutations(lbody, new Set([w]), wMut)
+  if (wMut.has(w)) return null
+  const ivMut = new Set()
+  findMutations([';', ...lbody.slice(1).filter((_, k) => k !== ivIdx - 1)], new Set([iv]), ivMut)
+  if (ivMut.has(iv)) return null
+  if (hasOuterContinue(lbody)) return null
+  if (cm.has(iv) || cm.has(w)) return null  // IV/divisor mutable via a closure call
+
+  const id = freshLoopId()
+  const cx = `__lsrx${id}`, cy = `__lsry${id}`
+  // seed (inside the w>0 branch): cx = (i%w)|0, cy = (i/w)|0 — one-time, i32 via |0
+  const seedDecls = [['=', cx, ['|', ['%', iv, w], 0]]]
+  if (usesDiv) seedDecls.push(['=', cy, ['|', ['/', iv, w], 0]])
+  const seed = ['let', ...seedDecls]
+
+  // transformed body: replace patterns, inject `cx++; if(cx>=w){cx=0; cy++}` after the increment
+  const newBody = [';']
+  for (let k = 1; k < lbody.length; k++) {
+    newBody.push(replace(lbody[k], iv, w, cx, cy))
+    if (k === ivIdx) {
+      newBody.push(['=', cx, ['+', cx, 1]])
+      if (usesDiv) newBody.push(['=', cy, ['+', cy, ['?:', ['>=', cx, w], 1, 0]]])
+      newBody.push(['=', cx, ['?:', ['>=', cx, w], 0, cx]])
+    }
+  }
+  const fast = ['while', replace(cond, iv, w, cx, cy), newBody]
+  // The counters track i%w only when w>0 AND i>=0: w<=0 gives NaN / a negative-divisor
+  // modulo, and i<0 makes i%w negative (JS modulo takes the dividend's sign) — neither
+  // follows the 0..w-1 +1-wrap the counters model. A one-time `w>0 && i>=0` guard (i is
+  // the IV's entry value; it only increments, so it stays ≥0) keeps the fast i32 path
+  // for the universal positive-dimension, forward-counting case and falls back to the
+  // unmodified loop otherwise — sound for any w, any start.
+  return [['if', ['&&', ['>', w, 0], ['>=', iv, 0]], ['{}', [';', seed, fast]], stmt]]
+}
+
+export function strengthReduceLoopDivMod(body) {
+  const cm = closureMutatedVars(body)
+  return rewriteBlocks(body, stmt => tryReduce(stmt, cm))
+}

package/src/compile/loop-model.js

@@ -0,0 +1,91 @@
+// Shared AST-level loop primitives for the per-function loop transforms
+// (loop-divmod, loop-square, peel-stencil). Operates on the post-prepare AST.
+//
+// Each of those passes recognizes one narrow loop idiom and rewrites it, but they had
+// re-derived the same building blocks verbatim: the `[, v]` number-literal-hole
+// recognizer, the "+1 induction variable" matcher, the closure-mutated-variable
+// analysis (a var assigned inside any `=>` may be mutated by a call in the loop, so it
+// is unsafe as an IV/bound/divisor), the while/for structural normalizer, and the
+// post-order block walk that applies a per-statement rewrite. One home for all of them
+// — adding the next loop transform is then one recognizer over these, not a fourth copy.
+
+import { ASSIGN_OPS } from '../ast.js'
+import { ctx } from '../ctx.js'
+
+// Fresh id for a loop transform's generated locals (`__lsrx<id>`, `__pks<id>`, …). Backed by
+// a per-compile counter (reset in ctx.reset), NOT a module-global: a module-`let` counter
+// grows unbounded across a long-lived host and makes compile(P) non-deterministic — its output
+// names depend on how many programs were compiled before it. Distinct prefixes keep the shared
+// id space collision-free across transforms.
+export const freshLoopId = () => ctx.transform.loopXformId++
+
+// Post-prepare number literals are sparse-array holes `[<hole>, v]` (length 2, the op
+// slot `n[0]` is the elided hole == null). `litVal` returns the numeric value or null;
+// `litN(n, k)` tests for the exact literal `k`.
+export const litVal = (n) => Array.isArray(n) && n.length === 2 && n[0] == null && typeof n[1] === 'number' ? n[1] : null
+export const litN = (n, k) => Array.isArray(n) && n.length === 2 && n[0] == null && n[1] === k
+
+// The induction variable a statement increments by exactly +1, else null. Covers
+// `i++` (post-inc desugars to `(++i) - 1`), `++i`, `i += 1`, `i = i + 1` / `i = 1 + i`.
+export function unitIncVar(stmt) {
+  if (!Array.isArray(stmt)) return null
+  let inc = stmt
+  if (stmt[0] === '-' && litN(stmt[2], 1) && Array.isArray(stmt[1]) && stmt[1][0] === '++') inc = stmt[1]
+  if (inc[0] === '++' && typeof inc[1] === 'string') return inc[1]
+  if (stmt[0] === '+=' && typeof stmt[1] === 'string' && litN(stmt[2], 1)) return stmt[1]
+  if (stmt[0] === '=' && typeof stmt[1] === 'string' && Array.isArray(stmt[2]) && stmt[2][0] === '+') {
+    const [, a, b] = stmt[2]
+    if (a === stmt[1] && litN(b, 1)) return stmt[1]
+    if (b === stmt[1] && litN(a, 1)) return stmt[1]
+  }
+  return null
+}
+
+// Normalize a `while` / flat `for` (`['for', init, cond, update, body]`) into a common
+// shape, else null. `init`/`step` are null for a while; a `for` with fewer than the 5
+// flat slots (no body) is not a loop here.
+export function normalizeLoop(stmt) {
+  if (!Array.isArray(stmt)) return null
+  if (stmt[0] === 'while') return { kind: 'while', init: null, cond: stmt[1], step: null, body: stmt[2] }
+  if (stmt[0] === 'for' && stmt.length >= 5) return { kind: 'for', init: stmt[1], cond: stmt[2], step: stmt[3], body: stmt[4] }
+  return null
+}
+
+// Variables assigned anywhere inside a closure (`=>`) in `body`. A call in the loop can
+// mutate such a var even though a direct-write scan (findMutations) misses it (the
+// closure may be defined outside the loop), so an IV / bound / divisor in this set is
+// unsafe to strength-reduce. (ASSIGN_OPS plus the `++`/`--` updates.)
+export function closureMutatedVars(body) {
+  const out = new Set()
+  const collect = (n) => {
+    if (!Array.isArray(n)) return
+    if (typeof n[1] === 'string' && (ASSIGN_OPS.has(n[0]) || n[0] === '++' || n[0] === '--')) out.add(n[1])
+    n.forEach(collect)
+  }
+  const walk = (n) => {
+    if (!Array.isArray(n)) return
+    if (n[0] === '=>') collect(n)
+    n.forEach(walk)
+  }
+  walk(body)
+  return out
+}
+
+// Post-order block rewriter: walk `body`; for every `;`-sequence, apply `tryStmt` to
+// each statement. A truthy result is an ARRAY of replacement statements spliced in
+// place; a falsy result keeps the statement unchanged. Children are rewritten first, so
+// a nested loop is transformed before the block that encloses it.
+export function rewriteBlocks(body, tryStmt) {
+  const walk = (node) => {
+    if (!Array.isArray(node)) return node
+    const n = node.map(walk)
+    if (n[0] !== ';') return n
+    const out = [';']
+    for (let k = 1; k < n.length; k++) {
+      const r = tryStmt(n[k])
+      if (r) out.push(...r); else out.push(n[k])
+    }
+    return out
+  }
+  return walk(body)
+}

package/src/compile/loop-square.js

@@ -0,0 +1,102 @@
+// Bounded-square narrowing: carry `i*i` as i32 inside a loop guarded by `i*i < CONST`.
+//
+// A PRODUCT `i*i` is f64 in unhinted JS — the integer-overflow contract: a product can
+// exceed 2³¹, where JS keeps a Number, so jz can't blindly use i32.mul (it would wrap).
+// In a Sieve-of-Eratosthenes `for(i=2; i*i<LIMIT; i++) for(j=i*i; j<LIMIT; j+=i) …`, that
+// makes the outer bound, the inner counter `j`'s init, and the whole index chain f64 —
+// each typed-array access then pays an f64→i32 convert.
+//
+// But when the loop GUARD is `i*i < CONST` with CONST a compile-time constant ≤ 2³⁰, the
+// counter is i ≤ ⌈√CONST⌉ ≤ 2¹⁵ throughout the body (the loop is still running, and i is
+// only incremented by +1), so EVERY `i*i` there is < 2³⁰ < 2³¹ and `Math.imul(i,i) == i*i`
+// exactly. The exit overshoot (the first i with i*i ≥ CONST, evaluated in the condition)
+// is ≤ CONST + 2√CONST+1 < 2³¹ for CONST ≤ 2³⁰, so the condition narrows soundly too. The
+// hard cap is 2³⁰ — well under the 2³¹ overflow point, with margin for the overshoot.
+//
+// So we rewrite those `i*i` to `Math.imul(i,i)` (jz's i32 multiply): semantically identical
+// in the proven range, and it lets the EXISTING i32 machinery carry the index chain as i32
+// — the inner `j` (init now i32, step `j+i` i32, bound `j<CONST` i32) cascades on its own.
+//
+// Sound iff: the loop condition is `(i*i) </≤ CONST` (CONST ≤ 2³⁰), and the IV `i` is
+// incremented by +1 and NOT otherwise mutated (so within a body iteration i ∈ {entry, +1},
+// both squares < 2³¹). Mirrors strength-reduce-divmod's structure (post-prepare `while`).
+
+import { findMutations } from './analyze-scans.js'
+import { includeMods } from '../autoload.js'
+import { ctx } from '../ctx.js'
+import { litVal, unitIncVar, normalizeLoop, closureMutatedVars, rewriteBlocks } from './loop-model.js'
+
+const SQUARE_BOUND_MAX = 2 ** 30
+// The constant numeric value of a bound: a literal, OR a module const folded to an int
+// (`const LIMIT = 1<<20` → ctx.scope.constInts.get('LIMIT') = 1048576 — the bench form).
+const boundVal = (n) => {
+  const lit = litVal(n)
+  if (lit != null) return lit
+  if (typeof n === 'string') { const v = ctx.scope.constInts?.get(n); return typeof v === 'number' ? v : null }
+  return null
+}
+// `i * i` — the IV squared.
+const isSquare = (n, iv) => Array.isArray(n) && n[0] === '*' && n[1] === iv && n[2] === iv && typeof iv === 'string'
+// Math.imul(i, i) in CANONICAL post-prepare form: prepare resolves `Math.imul` → the string
+// ref `'math.imul'`, so the call is `['()', 'math.imul', [',', i, i]]`. math.imul emits a
+// primitive `i32.mul` (no stdlib helper / module include needed).
+const imulOf = (iv) => ['()', 'math.imul', [',', iv, iv]]
+
+// The IV of a `(i*i) </≤ CONST` (or mirrored `CONST >/≥ (i*i)`) guard, CONST ≤ 2³⁰, else null.
+function boundedSquareIV(cond) {
+  if (!Array.isArray(cond)) return null
+  let prod, bound
+  if (cond[0] === '<' || cond[0] === '<=') { prod = cond[1]; bound = cond[2] }
+  else if (cond[0] === '>' || cond[0] === '>=') { prod = cond[2]; bound = cond[1] }
+  else return null
+  if (!isSquare(prod, prod && prod[1])) return null
+  const b = boundVal(bound)
+  if (b == null || b < 0 || b > SQUARE_BOUND_MAX) return null
+  return prod[1]
+}
+
+// Narrow a `for`/`while` whose guard is `(i*i) </≤ CONST` (CONST ≤ 2³⁰) and whose IV `i` is
+// incremented by +1 and not otherwise mutated — then within any body iteration i ∈ {entry,
+// entry+1}, entry ≤ ⌈√CONST⌉ ≤ 2¹⁵, so every `i*i` (and the exit overshoot) is < 2³¹ and
+// Math.imul(i,i) == i*i. Rewrites those products; the dependent counter chain cascades to i32.
+// `cm` is the function's closure-mutated-vars set (an IV in it has an unprovable entry value).
+function tryNarrow(stmt, cm) {
+  const L = normalizeLoop(stmt)
+  if (!L) return null
+  const { kind, cond, step, body } = L, isFor = kind === 'for'
+
+  const iv = boundedSquareIV(cond)
+  if (!iv) return null
+  if (cm.has(iv)) return null   // mutable via a closure call — entry value unprovable
+
+  if (isFor) {
+    // The `for` update is the IV's sole, +1 mutation; the body must not reassign i.
+    if (unitIncVar(step) !== iv) return null
+    const ivMut = new Set(); findMutations(body, new Set([iv]), ivMut)
+    if (ivMut.has(iv)) return null
+  } else {
+    // `while`: the increment lives in the body (exactly one +1; nothing else mutates i).
+    if (!Array.isArray(body) || body[0] !== ';') return null
+    let ivIdx = -1
+    for (let k = 1; k < body.length; k++) if (unitIncVar(body[k]) === iv) { if (ivIdx >= 0) return null; ivIdx = k }
+    if (ivIdx < 0) return null
+    const ivMut = new Set()
+    findMutations([';', ...body.slice(1).filter((_, k) => k !== ivIdx - 1)], new Set([iv]), ivMut)
+    if (ivMut.has(iv)) return null
+  }
+
+  // We're injecting `math.imul` after prepare's auto-import step, so ensure the math module
+  // (which registers the `math.imul` → i32.mul primitive emitter) is included.
+  includeMods('math')
+  // Rewrite every `i*i` in the condition + body to Math.imul(i,i) (NOT init/update — they're
+  // `i=2` / `i++`). The inner counter whose init this feeds cascades to i32 on its own.
+  const rw = (n) => !Array.isArray(n) ? n : isSquare(n, iv) ? imulOf(iv) : n.map(rw)
+  return [isFor
+    ? ['for', L.init, rw(cond), step, rw(body)]
+    : ['while', rw(cond), rw(body)]]
+}
+
+export function narrowBoundedSquare(body) {
+  const cm = closureMutatedVars(body)
+  return rewriteBlocks(body, stmt => tryNarrow(stmt, cm))
+}