jz 0.6.0 → 0.8.0

package/src/compile/plan/inline.js

@@ -39,12 +39,25 @@ import {
 // expression — null if void or no trailing return value.
 const inlinedBody = (func, args) => {
   const params = func.sig.params
-  if (args.length !== params.length || !args.every(isSimpleArg)) return null
+  if (args.length !== params.length) return null
   const paramNames = new Set(params.map(p => p.name))
   if (mutatesAny(func.body, paramNames)) return null
 
+  // A simple arg (ident / literal / arithmetic) is cheap to substitute directly, even when its
+  // param is used several times. A NON-simple arg (a call, `?:`, indexed load) is bound to a fresh
+  // temp evaluated ONCE in call order — preserving evaluation count + left-to-right order, and
+  // never duplicating the expression. This lets nested calls inline: `lerp(grad(a), grad(b), u)`
+  // binds `t0 = grad(a); t1 = grad(b)` and substitutes the body with t0/t1; a later inliner pass
+  // then folds grad into those temp decls (the fixpoint in inlineHotInternalCalls).
   const subst = new Map()
-  for (let i = 0; i < params.length; i++) subst.set(params[i].name, args[i])
+  const argPrefix = []
+  for (let i = 0; i < params.length; i++) {
+    const arg = args[i]
+    if (isSimpleArg(arg)) { subst.set(params[i].name, arg); continue }
+    const tmp = `${T}inarg${ctx.func.uniq++}`
+    argPrefix.push(['const', ['=', tmp, arg]])
+    subst.set(params[i].name, tmp)
+  }
 
   const locals = new Set()
   collectBindings(func.body, locals)
@@ -56,13 +69,13 @@ const inlinedBody = (func, args) => {
   const stmts = blockStmts(func.body)
   // Expression-bodied arrow `(c) => expr`: no statement block; the whole body
   // *is* the return value. Treat as zero-prefix + value.
-  if (!stmts) return { prefix: [], value: cloneWithSubst(func.body, subst, rename) }
+  if (!stmts) return { prefix: argPrefix, value: cloneWithSubst(func.body, subst, rename) }
   const last = stmts.length ? stmts[stmts.length - 1] : null
   const isTrailingReturn = Array.isArray(last) && last[0] === 'return'
   const prefixSrc = isTrailingReturn ? stmts.slice(0, -1) : stmts
   const prefix = prefixSrc.map(stmt => cloneWithSubst(stmt, subst, rename))
   const value = isTrailingReturn && last.length > 1 ? cloneWithSubst(last[1], subst, rename) : null
-  return { prefix, value }
+  return { prefix: argPrefix.length ? [...argPrefix, ...prefix] : prefix, value }
 }
 
 const stmtDeclName = (stmt) => {
@@ -109,7 +122,14 @@ const isCandidateCall = (node, candidates) =>
 // arithmetic>`, substituting the decls into the return value turns it into a
 // zero-prefix expression. Duplicated subtrees (`dx` used twice → `x1-x2`
 // twice) are pure, and the watr-layer CSE collapses the copies.
-const PURE_FLATTEN_OPS = new Set(['+', '-', '*', '/', '%', 'u-', 'u+', '&', '|', '^', '<<', '>>', '>>>'])
+const PURE_FLATTEN_OPS = new Set([
+  '+', '-', '*', '/', '%', 'u-', 'u+', '&', '|', '^', '<<', '>>', '>>>',
+  // pure value-producing ops with no side effects — safe to duplicate (CSE collapses copies):
+  // comparisons, logical, bit-not, and the conditional. Lets a branchy leaf like noise's
+  // `grad(h,x,y) => { …; u = (h&1)===0 ? x : -x; … }` flatten to an expression so it inlines
+  // into its multi-call caller `perlin` — `lerp(grad(a), grad(b), u)` then collapses end to end.
+  '<', '<=', '>', '>=', '==', '!=', '===', '!==', '&&', '||', '!', '~', '?:',
+])
 const pureFlattenExpr = (n) => {
   if (typeof n === 'number' || typeof n === 'string') return true  // literal or ident
   if (!Array.isArray(n)) return false
@@ -163,13 +183,19 @@ const inlineInExpr = (node, candidates) => {
 
 const inlineInStmt = (stmt, candidates, loopVariantNames = null) => {
   if (!Array.isArray(stmt)) return { node: stmt, changed: false }
-  // Statement-position call: discard return value, splice prefix in place.
+  // Statement-position call: the result is unused, but the callee's return
+  // EXPRESSION may still carry side effects — an expression-bodied arrow whose body
+  // is itself effectful (`seek = n => idx = n` inlines to the assignment `idx = i`;
+  // a one-liner that calls another fn) puts the effect in `value`, not `prefix`.
+  // Emit it as a trailing statement so the effect runs; a pure value is dropped
+  // later by vacuum/DCE. (Dropping it lost the parser's seek() idx-advance → ∞ loop.)
   if (isCandidateCall(stmt, candidates)) {
     const args = callArgs(stmt)
     const shape = args && inlinedBody(candidates.get(stmt[1]), args)
     if (shape) {
       const { hoisted, rest } = partitionInvariantPrefix(shape.prefix, loopVariantNames)
-      return { node: ['{}', [';', ...rest]], changed: true, splice: rest, hoisted }
+      const splice = shape.value !== null ? [...rest, shape.value] : rest
+      return { node: ['{}', [';', ...splice]], changed: true, splice, hoisted }
     }
   }
   // `let/const X = call(...)` with single decl: inline as prefix + decl(value).
@@ -259,9 +285,99 @@ const inlineInStmt = (stmt, candidates, loopVariantNames = null) => {
   return { node: stmt, changed: false }
 }
 
+// Short-circuit operators: only the FIRST operand is unconditionally evaluated; a call in
+// a later operand might not run, so it can't be hoisted.
+const SHORT_CIRCUIT = new Set(['?:', '?', '&&', '||', '??'])
+// Optional chaining: jz's own desugaring already tees the base to evaluate it once, and
+// the key/args run conditionally — so the hoist treats the WHOLE expression as opaque (no
+// operand, not even the base, is hoisted out) to avoid colliding with that desugaring.
+const OPTIONAL_CHAIN = new Set(['?.', '?.[]', '?.()'])
+// Mutating expression operators — evaluating one is an observable side effect.
+const ASSIGN_OPS = new Set(['=', '+=', '-=', '*=', '/=', '%=', '&=', '|=', '^=', '<<=', '>>=', '>>>=', '**=', '&&=', '||=', '??=', '++', '--'])
+// Does evaluating this expression have an observable side effect (a call or assignment)?
+const containsEffect = (n) => Array.isArray(n) && n[0] !== '=>' &&
+  (n[0] === '()' || n[0] === '?.()' || ASSIGN_OPS.has(n[0]) || n.slice(1).some(containsEffect))
+
+// Hoist an unconditionally-evaluated NESTED call to a block-body candidate out to a
+// preceding `const __h = call(...)` temp. inlineInStmt folds block-body candidates only at
+// a DIRECT `const X = call` / `X = call`; a call buried in an expression (noise's
+// `sum = sum + amp * perlin(x)`) is reached by neither that path nor inlineInExpr (which
+// only substitutes zero-prefix expr-bodies). Hoisting normalizes it to the direct form.
+// Only block-body candidates and only unconditional positions — preserving evaluation order
+// + count. Statement HEADERS that are expression positions (for-init/update, while/if test)
+// are left untouched: there's no sound place for a hoisted decl there, so those calls just
+// stay outlined. Conservatively leaves unrecognized statement shapes alone.
+const hoistNestedCalls = (body, blockNames) => {
+  if (!blockNames.size || !Array.isArray(body)) return { node: body, changed: false }
+  let changed = false
+  const seq = (stmts) => stmts.length === 1 ? stmts[0] : [';', ...stmts]
+  // Lifting a call to the pre-decl block moves its evaluation to the TOP of the statement.
+  // Sound only if no observable side effect is evaluated BEFORE it — else its effect jumps
+  // ahead of that one (`a() + helper(x)` must keep a()'s effect first). `eff.seen` threads
+  // left-to-right through evaluation order: a call or assignment LEFT IN PLACE marks every
+  // later position. A hoisted call moves as a unit — its args run in a fresh inner eff, and
+  // it does NOT advance the outer eff (the whole unit relocates together, order intact).
+  const hExpr = (n, pre, cond, eff) => {
+    if (!Array.isArray(n) || n[0] === '=>') return n
+    if (!cond && !eff.seen && n[0] === '()' && typeof n[1] === 'string' && blockNames.has(n[1])) {
+      const call = [n[0], n[1], ...n.slice(2).map(a => hExpr(a, pre, false, { seen: false }))]
+      const tmp = `${T}inl${ctx.func.uniq++}_h`
+      pre.push(['const', ['=', tmp, call]])
+      changed = true
+      return [null, tmp]
+    }
+    if (OPTIONAL_CHAIN.has(n[0])) {
+      const out = [n[0], ...n.slice(1).map(c => hExpr(c, pre, true, eff))]
+      if (n[0] === '?.()') eff.seen = true  // optional CALL may run
+      return out
+    }
+    if (SHORT_CIRCUIT.has(n[0]))
+      return [n[0], hExpr(n[1], pre, cond, eff), ...n.slice(2).map(c => hExpr(c, pre, true, eff))]
+    const out = [n[0], ...n.slice(1).map(c => hExpr(c, pre, cond, eff))]
+    if (n[0] === '()' || ASSIGN_OPS.has(n[0])) eff.seen = true  // a call/assign left in place is an effect
+    return out
+  }
+  // A RHS that is DIRECTLY a candidate call is already folded by inlineInStmt's
+  // `const X = call` / `X = call` paths — hoisting it would be redundant and (for an
+  // object/array-literal `{}`-bodied factory) would break the post-inline alias chain.
+  // Only hoist NESTED calls; leave a top-level direct call to those paths.
+  const directCall = (e) => Array.isArray(e) && e[0] === '()' && typeof e[1] === 'string' && blockNames.has(e[1])
+  const hStmt = (s) => {  // → array of statements (hoisted decls prepended)
+    if (!Array.isArray(s)) return [s]
+    switch (s[0]) {
+      case ';': return [[';', ...s.slice(1).flatMap(hStmt)]]
+      case '{}': return [['{}', seq(hStmt(s[1]))]]
+      case 'if': return [s.length > 3
+        ? ['if', s[1], seq(hStmt(s[2])), seq(hStmt(s[3]))]
+        : ['if', s[1], seq(hStmt(s[2]))]]
+      case 'for': { const i = forLoopBodyIndex(s); return [withForLoopBody(s, seq(hStmt(s[i])))] }
+      case 'while': return [['while', s[1], seq(hStmt(s[2]))]]
+      case 'let': case 'const': {
+        if (s.length === 2 && Array.isArray(s[1]) && s[1][0] === '=' && typeof s[1][1] === 'string' && !directCall(s[1][2])) {
+          const pre = []; const rhs = hExpr(s[1][2], pre, false, { seen: false })
+          return pre.length ? [...pre, [s[0], ['=', s[1][1], rhs]]] : [s]
+        }
+        return [s]
+      }
+      // A computed assign target (`a[i]=…`) evaluates its index BEFORE the RHS, so an effect
+      // there (`a[j++]=…`) must block hoisting too — seed eff.seen from the LHS.
+      case '=': { if (directCall(s[2])) return [s]; const pre = []; const rhs = hExpr(s[2], pre, false, { seen: containsEffect(s[1]) }); return pre.length ? [...pre, ['=', s[1], rhs]] : [s] }
+      case 'return': { if (s.length < 2 || directCall(s[1])) return [s]; const pre = []; const v = hExpr(s[1], pre, false, { seen: false }); return pre.length ? [...pre, ['return', v]] : [s] }
+      default: return [s]  // unrecognized shape (break/continue/throw/try/switch): leave alone
+    }
+  }
+  const out = hStmt(body)
+  return { node: changed ? seq(out) : body, changed }
+}
+
 export const inlineHotInternalCalls = (programFacts, ast) => {
   const cfg = ctx.transform.optimize
   if (cfg && cfg.sourceInline === false) return false
+  // Transitive candidacy + expression-position hoisting are a size↔speed trade (they
+  // pull a large multi-call leaf like noise's perlin fully into its hot caller, where
+  // the lower tiers prefer to keep multi-caller helpers outlined for V8 tier-up). Gate
+  // both on the speed tier so levels ≤2 keep their conservative inlining policy.
+  const speedTier = !!(cfg && cfg.inlineFns)
 
   const fixedByFunc = new Map(ctx.func.list.map(func => [func, fixedTypedArraysInBody(func.body)]))
   const typedByFunc = new Map(ctx.func.list.map(func => [func, analyzeBody(func.body).typedElems]))
@@ -315,7 +431,14 @@ export const inlineHotInternalCalls = (programFacts, ast) => {
   // entry points, not pulling a leaf INTO an export's hot loop (game-of-life's
   // step calls rot per cell; pre-Turboshaft wasm tiers never inline calls).
   const leaves = new Set()
+  // Transitive candidacy via fixpoint: a function whose only user-callees are THEMSELVES
+  // candidates (so they inline away) can be inlined too. noise's `perlin` calls grad/fade/
+  // lerp (loop-free leaves) — once those are candidates, perlin clears the call-bearing-body
+  // gate and becomes a leaf candidate. Each pass adds ≥1 or stops, so it's bounded.
+  for (let recollect = true; recollect;) {
+  recollect = false
   for (const func of ctx.func.list) {
+    if (candidates.has(func.name)) continue
     const sites = sitesByCallee.get(func.name)
     // Exported leaf/kernel with exactly one internal caller (e.g. fill→beat in
     // floatbeat): inline into the caller's loop but keep the export for external
@@ -336,7 +459,13 @@ export const inlineHotInternalCalls = (programFacts, ast) => {
     const fullyFixedTypedArraySite = hasFullyFixedTypedArraySites(func, sites)
     const hasLoop = some(func.body, n => LOOP_OPS.has(n[0]))
     const isTinyLeaf = !hasLoop && nodeSize(func.body) <= 15
-    if (!sites || sites.length < 1 || (!isTinyLeaf && !fixedTypedArraySite && sites.length > 2) || sites.length > 8) continue
+    // A small leaf (no loop, ≤40 nodes) is cheap to splice even when called several times — its
+    // per-call overhead + lost cross-call fusion dwarfs the ≤8× duplication, and temp-binding +
+    // flattenPrefix keep the spliced body bounded (no arg re-evaluation, CSE collapses copies).
+    // The 2-site non-tiny-leaf cap would otherwise outline a hot helper like noise's `grad`
+    // (~30 nodes, called 4× from perlin) and freeze the call overhead per pixel.
+    const isSmallLeaf = !hasLoop && nodeSize(func.body) <= 48
+    if (!sites || sites.length < 1 || (!isTinyLeaf && !isSmallLeaf && !fixedTypedArraySite && sites.length > 2) || sites.length > 8) continue
     const stmts = blockStmts(func.body)
     // Expression-bodied arrow funcs (`(c) => expr`) have no block — body IS the
     // return value. Treat as a "tiny leaf" branch handled below; force hasLoop=false.
@@ -356,15 +485,24 @@ export const inlineHotInternalCalls = (programFacts, ast) => {
     // that get hammered from a hot caller's loop — replacing the call with its
     // body saves the per-iteration call+reinterpret overhead (tokenizer hot path).
     if (!hasLoop) {
-      if (some(func.body, n => n[0] === '()' && typeof n[1] === 'string' && ctx.func.names.has(n[1]))) continue
+      // Calls to functions that are THEMSELVES candidates are fine — they inline away;
+      // only a call to a non-candidate user function blocks (a later fixpoint pass re-checks).
+      // Speed-tier only; lower tiers keep the strict "any user call ⇒ outline" rule.
+      if (some(func.body, n => n[0] === '()' && typeof n[1] === 'string' && ctx.func.names.has(n[1]) && !(speedTier && candidates.has(n[1])))) continue
       // Per-iteration call overhead dwarfs body-size bloat when EVERY site sits
       // inside a caller's loop (game-of-life's rot: ~40 nodes × 2 sites, fired
-      // for most of 260k cells/frame). V8's pre-Turboshaft wasm tiers never
+      // for most of 260k cells/frame; cloth's relax: ~160 nodes × 2 sites, fired
+      // per link every relaxation pass). V8's pre-Turboshaft wasm tiers never
       // inline cross-function, so an out-of-line leaf in a hot loop is a hard
       // per-cell tax on Node ≤ 22 — and still saves call setup on newer tiers.
+      // The in-loop cap is generous because the gate above bounds non-tiny leaves
+      // to ≤2 sites, so the spliced duplication is at most ~2× a bounded body.
       const allSitesInLoop = sites.every(site =>
         site.callerFunc?.body && containsNode(site.callerFunc.body, site.node, false))
-      if (nodeSize(func.body) > (allSitesInLoop ? 80 : 30)) continue
+      // Non-in-loop cap is 40 (not 30) so a small leaf called from a straight-line but
+      // transitively-hot caller still inlines (noise's grad is called from perlin, which has
+      // no loop of its own but is itself the per-pixel kernel). Still tightly bounded.
+      if (nodeSize(func.body) > (allSitesInLoop ? 200 : 48)) continue
     }
     if (some(func.body, n => n[0] === '()' && n[1] === func.name)) continue
     // Kernels with nested loops (depth ≥ 2) are typically large and the inner
@@ -383,6 +521,8 @@ export const inlineHotInternalCalls = (programFacts, ast) => {
       forwarders.add(func.name)
     if (!hasLoop) leaves.add(func.name)
     candidates.set(func.name, func)
+    if (speedTier) recollect = true  // only the speed-tier transitive relaxation needs a re-pass
+  }
   }
   if (!candidates.size) return false
 
@@ -433,20 +573,38 @@ export const inlineHotInternalCalls = (programFacts, ast) => {
     // Route those through inlineInExpr so the call is replaced by the inlined
     // value expression instead.
     const isExprBody = !Array.isArray(func.body) || func.body[0] !== '{}'
-    const r = isExprBody
-      ? inlineInExpr(func.body, activeCandidates)
-      : inlineInStmt(func.body, activeCandidates)
-    let body = r.changed ? r.node : func.body
-    let bodyChanged = r.changed
-    // Expression-position pass. Exported callers take the leaf-safe subset —
-    // the same tier-up rationale as the statement path (leaves into exports
-    // are fine; relocated kernels are not).
+    // Expression-position pass takes the leaf-safe subset for exports — the same tier-up
+    // rationale as the statement path (leaves into exports are fine; relocated kernels are not).
     const exprActive = func.exported
       ? new Map([...exprOnlyCandidates].filter(([n]) => exportedCandidates.has(n)))
       : exprOnlyCandidates
-    if (exprActive.size) {
-      const e = inlineInExpr(body, exprActive)
-      if (e.changed) { body = e.node; bodyChanged = true }
+    // Iterate to a (bounded) fixpoint: inlining a call whose args are themselves candidate calls
+    // binds those args to temps (`t0 = grad(a)`); the next pass folds the candidate into the temp
+    // decl. Depth is bounded by call nesting (a small constant), capped here so a pathological
+    // chain can't loop unbounded.
+    // Loop-free block-body LEAVES: the stmt path folds them only at a DIRECT `const X =
+    // call`, never nested in an expression. Hoisting such a call to a temp (in the iter
+    // fixpoint below) lets inlineInStmt then fold it — the noise `sum + amp*perlin(x)`
+    // shape. Restricted to LEAVES (no own loop): a loop kernel called in expression
+    // position (e.g. a 2-site `reduce`) was deliberately staying outlined for V8 tier-up,
+    // and hoisting it would pull the loop into a cold caller.
+    const blockNames = new Set()
+    for (const n of activeCandidates.keys()) if (leaves.has(n) && !exprActive.has(n)) blockNames.add(n)
+    let body = func.body, bodyChanged = false
+    for (let iter = 0; iter < 4; iter++) {
+      let iterChanged = false
+      if (speedTier && !isExprBody && blockNames.size) {
+        const h = hoistNestedCalls(body, blockNames)
+        if (h.changed) { body = h.node; iterChanged = true }
+      }
+      const r = isExprBody ? inlineInExpr(body, activeCandidates) : inlineInStmt(body, activeCandidates)
+      if (r.changed) { body = r.node; iterChanged = true }
+      if (exprActive.size) {
+        const e = inlineInExpr(body, exprActive)
+        if (e.changed) { body = e.node; iterChanged = true }
+      }
+      if (!iterChanged) break
+      bodyChanged = true
     }
     if (bodyChanged) { func.body = body; changed = true }
   }