@aardworx/wombat.rendering 0.8.0 → 0.9.0

package/dist/runtime/heapEffectIR.js

@@ -0,0 +1,1161 @@
+// heapEffectIR — IR-level rewriter for the heap render path.
+//
+// Takes a wombat.shader Effect, transforms its IR Module so that:
+//   - Vertex attribute reads are replaced with heap-buffer pulls
+//     (cyclic via `vid % length`).
+//   - Uniform reads (in any stage) are replaced with heap-buffer
+//     decodes via per-draw refs from the headers buffer.
+//   - FS uniform decodes thread the per-uniform refs through VsOut
+//     as flat-interpolated `_h_<name>Ref: u32` varyings.
+//   - The heap arena's four typed views (heapU32, headersU32, heapF32,
+//     heapV4f) are added as storage-buffer bindings.
+//
+// The result is a fresh Module that compiles via the existing
+// wombat.shader pipeline (DCE, cross-stage link, fragment-output
+// linker, WGSL emit) — no string-regex post-processing.
+//
+// This module knows nothing about the runtime; the BucketLayout it
+// consumes is identical to the one heapScene.ts already builds.
+import { compileModule, stage as makeStage, effect as makeEffect } from "@aardworx/wombat.shader";
+import { substituteInputsInStage, readInputs, mapExpr, mapStmt, liftReturns } from "@aardworx/wombat.shader/passes";
+import { Tu32, Tf32, Vec, } from "@aardworx/wombat.shader/ir";
+import { megacallSearchPrelude, atlasVaryingNames, generateAtlasBindings, generateAtlasSwitch, generateAtlasPrelude, } from "./heapEffect.js";
+// ─── IR builders ────────────────────────────────────────────────────
+const Tbool = { kind: "Bool" };
+const constU32 = (n) => ({
+    kind: "Const", type: Tu32, value: { kind: "Int", signed: false, value: n },
+});
+const constF32 = (n) => ({
+    kind: "Const", type: Tf32, value: { kind: "Float", value: n },
+});
+const readScope = (scope, name, type) => ({
+    kind: "ReadInput", scope, name, type,
+});
+const item = (target, index, elemType) => ({
+    kind: "Item", target, index, type: elemType,
+});
+const add = (lhs, rhs, type) => ({ kind: "Add", lhs, rhs, type });
+const mul = (lhs, rhs, type) => ({ kind: "Mul", lhs, rhs, type });
+const div = (lhs, rhs, type) => ({ kind: "Div", lhs, rhs, type });
+const mod = (lhs, rhs, type) => ({ kind: "Mod", lhs, rhs, type });
+const eqU32 = (lhs, rhs) => ({ kind: "Eq", lhs, rhs, type: Tbool });
+const select = (cond, ifTrue, ifFalse, type) => ({
+    kind: "Conditional", cond, ifTrue, ifFalse, type,
+});
+const newVec = (components, type) => ({
+    kind: "NewVector", components, type,
+});
+const matFromRows = (rows, type) => ({
+    kind: "MatrixFromRows", rows, type,
+});
+// Runtime-sized array types for the heap storage buffers.
+const arrU32 = { kind: "Array", element: Tu32, length: "runtime" };
+const arrF32 = { kind: "Array", element: Tf32, length: "runtime" };
+const arrVec4 = { kind: "Array", element: Vec(Tf32, 4), length: "runtime" };
+const Tvec2 = Vec(Tf32, 2);
+const Tvec3 = Vec(Tf32, 3);
+const Tvec4 = Vec(Tf32, 4);
+// ─── Storage-buffer bindings ────────────────────────────────────────
+//
+// Bindings 0..3 are the four typed views of the heap arena. Bindings
+// 4+ (textures/samplers) are added by the layout / runtime as needed.
+function heapStorageBufferDecls() {
+    return [
+        { kind: "StorageBuffer", binding: { group: 0, slot: 0 }, name: "heapU32", layout: arrU32, access: "read" },
+        { kind: "StorageBuffer", binding: { group: 0, slot: 1 }, name: "headersU32", layout: arrU32, access: "read" },
+        { kind: "StorageBuffer", binding: { group: 0, slot: 2 }, name: "heapF32", layout: arrF32, access: "read" },
+        { kind: "StorageBuffer", binding: { group: 0, slot: 3 }, name: "heapV4f", layout: arrVec4, access: "read" },
+    ];
+}
+// ─── Heap-load IR builders ─────────────────────────────────────────
+//
+// Storage buffers are exposed in IR as `ReadInput("Uniform", name)`
+// (the same scope WGSL uses for all bindings). Indexing is `Item`.
+// Address arithmetic is plain Add/Mul on u32.
+const heapU32 = readScope("Uniform", "heapU32", arrU32);
+const headersU32 = readScope("Uniform", "headersU32", arrU32);
+const heapF32 = readScope("Uniform", "heapF32", arrF32);
+const heapV4f = readScope("Uniform", "heapV4f", arrVec4);
+/** `headersU32[drawIdx * stride + offset]`. */
+function loadHeaderRef(drawIdx, fieldByteOffset, strideU32) {
+    // u32 index = drawIdx * strideU32 + (fieldByteOffset / 4).
+    const offU32 = fieldByteOffset / 4;
+    const idx = add(mul(drawIdx, constU32(strideU32), Tu32), constU32(offU32), Tu32);
+    return item(headersU32, idx, Tu32);
+}
+/**
+ * Cyclic per-element index: `(elemIdx % count) * floatsPerElement`,
+ * where count = alloc-header length (u32 at offset 4, i.e. ref/4 + 1).
+ * For per-vertex attributes elemIdx = vid; for instance arrays =
+ * iidx (length is then the instance count, no cycling needed but the
+ * mod is harmless).
+ */
+function cyclicElemIdx(refIdent, elemIdx, floatsPerElem) {
+    const count = item(heapU32, add(div(refIdent, constU32(4), Tu32), constU32(1), Tu32), Tu32);
+    const cycled = mod(elemIdx, count, Tu32);
+    return mul(cycled, constU32(floatsPerElem), Tu32);
+}
+/**
+ * Load a per-draw uniform value from the heap given its `u32` ref.
+ * Type-driven: matrix → 4 vec4 columns from heapV4f; vec4 → one
+ * heapV4f read; vec3/2/scalar → heapF32 reads + NewVector.
+ */
+function loadUniformByRef(refIdent, wgslType) {
+    switch (wgslType) {
+        case "mat4x4<f32>": {
+            // base vec4 index = (ref + 16) / 16
+            const base = div(add(refIdent, constU32(16), Tu32), constU32(16), Tu32);
+            const cols = [];
+            for (let i = 0; i < 4; i++) {
+                const idx = i === 0 ? base : add(base, constU32(i), Tu32);
+                cols.push(item(heapV4f, idx, Tvec4));
+            }
+            return matFromRows(cols, { kind: "Matrix", element: Tf32, rows: 4, cols: 4 });
+        }
+        case "vec4<f32>":
+            return item(heapV4f, div(add(refIdent, constU32(16), Tu32), constU32(16), Tu32), Tvec4);
+        case "vec3<f32>": {
+            const base = div(add(refIdent, constU32(16), Tu32), constU32(4), Tu32);
+            return newVec([
+                item(heapF32, base, Tf32),
+                item(heapF32, add(base, constU32(1), Tu32), Tf32),
+                item(heapF32, add(base, constU32(2), Tu32), Tf32),
+            ], Tvec3);
+        }
+        case "vec2<f32>": {
+            const base = div(add(refIdent, constU32(16), Tu32), constU32(4), Tu32);
+            return newVec([
+                item(heapF32, base, Tf32),
+                item(heapF32, add(base, constU32(1), Tu32), Tf32),
+            ], Tvec2);
+        }
+        case "f32":
+            return item(heapF32, div(add(refIdent, constU32(16), Tu32), constU32(4), Tu32), Tf32);
+        case "u32":
+            return item(heapU32, div(add(refIdent, constU32(16), Tu32), constU32(4), Tu32), Tu32);
+        default:
+            throw new Error(`heapEffectIR: no IR loader for uniform type '${wgslType}'`);
+    }
+}
+/**
+ * Load an attribute value with cyclic addressing (`(idx % length) *
+ * stride`). The mod lets a length-1 broadcast replicate across all
+ * vertices. For vec4 the source can be V3-tight (12 B, .w = 1.0) or
+ * V4-tight (16 B, .w from data); we read the actual stride from the
+ * alloc header at offset 8, then use `select` for the .w fill.
+ */
+function loadAttributeByRef(refIdent, idx, wgslType) {
+    const dataF32Base = div(add(refIdent, constU32(16), Tu32), constU32(4), Tu32);
+    switch (wgslType) {
+        case "vec3<f32>": {
+            const off = cyclicElemIdx(refIdent, idx, 3);
+            const base = add(dataF32Base, off, Tu32);
+            return newVec([
+                item(heapF32, base, Tf32),
+                item(heapF32, add(base, constU32(1), Tu32), Tf32),
+                item(heapF32, add(base, constU32(2), Tu32), Tf32),
+            ], Tvec3);
+        }
+        case "vec4<f32>": {
+            // strideF = stride_bytes / 4  (header u32 at ref/4 + 2).
+            const strideBytes = item(heapU32, add(div(refIdent, constU32(4), Tu32), constU32(2), Tu32), Tu32);
+            const strideF = div(strideBytes, constU32(4), Tu32);
+            const cycled = mod(idx, item(heapU32, add(div(refIdent, constU32(4), Tu32), constU32(1), Tu32), Tu32), Tu32);
+            const off = mul(cycled, strideF, Tu32);
+            const base = add(dataF32Base, off, Tu32);
+            const w = select(eqU32(strideF, constU32(4)), item(heapF32, add(base, constU32(3), Tu32), Tf32), constF32(1.0), Tf32);
+            return newVec([
+                item(heapF32, base, Tf32),
+                item(heapF32, add(base, constU32(1), Tu32), Tf32),
+                item(heapF32, add(base, constU32(2), Tu32), Tf32),
+                w,
+            ], Tvec4);
+        }
+        case "vec2<f32>": {
+            const off = cyclicElemIdx(refIdent, idx, 2);
+            const base = add(dataF32Base, off, Tu32);
+            return newVec([
+                item(heapF32, base, Tf32),
+                item(heapF32, add(base, constU32(1), Tu32), Tf32),
+            ], Tvec2);
+        }
+        default:
+            throw new Error(`heapEffectIR: no IR loader for attribute type '${wgslType}'`);
+    }
+}
+// ─── Instance-index / vertex-index built-in plumbing ────────────────
+//
+// The user's stages don't have @builtin(instance_index)/(vertex_index)
+// inputs — they read inputs by name (vertexAttributes, uniforms).
+// We need those builtins to address the heap. Inject them as
+// EntryParameters on every vertex stage.
+function entryHasBuiltin(e, semantic) {
+    return e.inputs.some(p => p.decorations.some(d => d.kind === "Builtin" && d.value === semantic));
+}
+function builtinParam(semantic) {
+    return {
+        name: semantic,
+        type: Tu32,
+        semantic,
+        decorations: [{ kind: "Builtin", value: semantic }],
+    };
+}
+function injectVsBuiltins(m) {
+    const values = m.values.map((v) => {
+        if (v.kind !== "Entry" || v.entry.stage !== "vertex")
+            return v;
+        let inputs = v.entry.inputs;
+        if (!entryHasBuiltin(v.entry, "vertex_index")) {
+            inputs = [...inputs, builtinParam("vertex_index")];
+        }
+        if (!entryHasBuiltin(v.entry, "instance_index")) {
+            inputs = [...inputs, builtinParam("instance_index")];
+        }
+        return { ...v, entry: { ...v.entry, inputs } };
+    });
+    return { ...m, values };
+}
+// ─── VS body rewrites ──────────────────────────────────────────────
+//
+// For each surviving uniform field, build a heap-load expression
+// using `instance_index` as drawIdx. For each surviving attribute,
+// use `vertex_index` and cyclic addressing.
+//
+// These map builders extract the (name → Expr) mapping that the VS
+// substitution needs; the actual rewrite is delegated to
+// `Effect.substitute({ vertex: { uniforms, inputs } })` in
+// `compileHeapEffectIR`.
+/**
+ * Per-RO selector expression. Always reads the `heap_drawIdx` local
+ * defined by the megacall prelude.
+ */
+function drawIdxExprFor(_layout) {
+    return { kind: "Var", var: { name: "heap_drawIdx", type: Tu32, mutable: false } };
+}
+/**
+ * `(uniformName → loadUniformByRef(loadHeaderRef(...), wgslType))` for
+ * every uniform-ref field in `layout`. Used as the `vertex.uniforms`
+ * argument to `Effect.substitute`.
+ */
+function buildVertexUniformMap(layout) {
+    const drawIdx = drawIdxExprFor(layout);
+    const iidx = readScope("Builtin", "instance_index", Tu32);
+    const stride = layout.strideU32;
+    const out = new Map();
+    for (const f of layout.drawHeaderFields) {
+        if (f.kind !== "uniform-ref")
+            continue;
+        const refExpr = loadHeaderRef(drawIdx, f.byteOffset, stride);
+        const value = layout.perInstanceUniforms.has(f.name)
+            ? loadInstanceByRef(refExpr, iidx, f.uniformWgslType ?? "")
+            : loadUniformByRef(refExpr, f.uniformWgslType ?? "");
+        out.set(f.name, value);
+    }
+    return out;
+}
+/**
+ * `(attributeName → loadAttributeByRef(loadHeaderRef(...), idx, wgslType))`
+ * for every attribute-ref field in `layout`. Used as the
+ * `vertex.inputs` argument to `Effect.substitute`.
+ */
+function buildVertexAttributeMap(layout) {
+    const drawIdx = drawIdxExprFor(layout);
+    const iidx = readScope("Builtin", "instance_index", Tu32);
+    const vid = readScope("Builtin", "vertex_index", Tu32);
+    const stride = layout.strideU32;
+    const out = new Map();
+    for (const f of layout.drawHeaderFields) {
+        if (f.kind !== "attribute-ref")
+            continue;
+        const refExpr = loadHeaderRef(drawIdx, f.byteOffset, stride);
+        const value = layout.perInstanceAttributes.has(f.name)
+            ? loadAttributeByRef(refExpr, iidx, f.attributeWgslType ?? "")
+            : loadAttributeByRef(refExpr, vid, f.attributeWgslType ?? "");
+        out.set(f.name, value);
+    }
+    return out;
+}
+function mapVertexEntryBodies(m, fn) {
+    return mapStageEntryBodies(m, "vertex", fn);
+}
+function mapStageEntryBodies(m, stage, fn) {
+    const values = m.values.map((v) => {
+        if (v.kind !== "Entry" || v.entry.stage !== stage)
+            return v;
+        return { ...v, entry: { ...v.entry, body: fn(v.entry.body) } };
+    });
+    return { ...m, values };
+}
+/** Append `extras` to every VS entry's outputs and every FS entry's inputs. */
+function addInterstageParams(m, extras) {
+    if (extras.length === 0)
+        return m;
+    const values = m.values.map((v) => {
+        if (v.kind !== "Entry")
+            return v;
+        if (v.entry.stage === "vertex") {
+            return { ...v, entry: { ...v.entry, outputs: [...v.entry.outputs, ...extras] } };
+        }
+        if (v.entry.stage === "fragment") {
+            return { ...v, entry: { ...v.entry, inputs: [...v.entry.inputs, ...extras] } };
+        }
+        return v;
+    });
+    return { ...m, values };
+}
+/**
+ * Prepend `stmts` to every VS entry's body. Flatten when the body is
+ * already a `Sequential` — CSE numbers `_cse0..N` per Sequential, so
+ * nested Sequentials produce colliding declarations in the emitted
+ * WGSL.
+ */
+function prependVsBodyStmts(m, stmts) {
+    if (stmts.length === 0)
+        return m;
+    return mapVertexEntryBodies(m, body => {
+        if (body.kind === "Sequential") {
+            return { kind: "Sequential", body: [...stmts, ...body.body] };
+        }
+        return { kind: "Sequential", body: [...stmts, body] };
+    });
+}
+// ─── FS uniform threading ──────────────────────────────────────────
+//
+// FS can't access `instance_index` directly, so we thread the per-
+// uniform refs through VsOut as flat-interpolated u32 varyings:
+//
+//   - For each FS-used uniform `X` whose layout entry is uniform-ref:
+//       VS output: `_h_<X>Ref: u32 @interpolate(flat)`
+//       VS body:   `WriteOutput("_h_<X>Ref", loadHeaderRef(drawIdx, X))`
+//       FS input:  `_h_<X>Ref: u32 @interpolate(flat)`
+//       FS body:   `ReadInput("Uniform", X)` →
+//                  `loadUniformByRef(ReadInput("Input", "_h_<X>Ref"), …)`
+//
+// For per-instance uniforms (layout.perInstanceUniforms), iidx is also
+// threaded as `_iidx: u32 @interpolate(flat)` and the FS uses
+// `loadInstanceByRef` instead.
+function wgslTypeFromHeapField(wgslType) {
+    switch (wgslType) {
+        case "mat4x4<f32>": return { kind: "Matrix", element: Tf32, rows: 4, cols: 4 };
+        case "vec4<f32>": return Tvec4;
+        case "vec3<f32>": return Tvec3;
+        case "vec2<f32>": return Tvec2;
+        case "f32": return Tf32;
+        case "u32": return Tu32;
+        default: throw new Error(`heapEffectIR: unknown uniform wgsl type '${wgslType}'`);
+    }
+}
+/**
+ * Family-member FS uniform rewrite (direct heap reads).
+ *
+ * The wrapper module threads `heap_drawIdx` and `instId` as plain u32
+ * flat varyings into the FS, then passes them to each per-effect FS
+ * helper. Inside the helper, every `ReadInput("Uniform", X)` is
+ * substituted with `loadUniformByRef(loadHeaderRef(heap_drawIdx, …), …)`
+ * (or `loadInstanceByRef(refExpr, instId, …)` for per-instance
+ * uniforms). No VS-side WriteOutput stmts and no `_h_*Ref` varyings —
+ * the carrier is the single `heap_drawIdx` u32.
+ *
+ * `heap_drawIdx` and `instId` are referenced as `Var` exprs; the post-
+ * emit `applyFamilyMemberFsShape` rewires the FS entry signature to
+ * surface them as plain u32 fn parameters.
+ */
+function rewriteFsUniformsDirect(m, layout) {
+    const used = new Set();
+    for (const v of m.values) {
+        if (v.kind !== "Entry" || v.entry.stage !== "fragment")
+            continue;
+        for (const sn of readInputs(v.entry.body).values()) {
+            if (sn.scope === "Uniform")
+                used.add(sn.name);
+        }
+    }
+    if (used.size === 0)
+        return m;
+    const fieldByName = new Map(layout.drawHeaderFields.map(f => [f.name, f]));
+    const drawIdxExpr = { kind: "Var", var: { name: "heap_drawIdx", type: Tu32, mutable: false } };
+    const instIdExpr = { kind: "Var", var: { name: "instId", type: Tu32, mutable: false } };
+    const fsSubst = new Map();
+    for (const name of used) {
+        const f = fieldByName.get(name);
+        if (f === undefined || f.kind !== "uniform-ref")
+            continue;
+        const refExpr = loadHeaderRef(drawIdxExpr, f.byteOffset, layout.strideU32);
+        const value = layout.perInstanceUniforms.has(name)
+            ? loadInstanceByRef(refExpr, instIdExpr, f.uniformWgslType ?? "")
+            : loadUniformByRef(refExpr, f.uniformWgslType ?? "");
+        fsSubst.set(name, value);
+    }
+    return substituteInputsInStage(m, "fragment", "Uniform", n => fsSubst.get(n));
+}
+/**
+ * Family-member FS atlas rewrite (direct heap reads).
+ *
+ * Like `rewriteFsAtlasTextures` but instead of threading
+ * `_h_<name>{PageRef,FormatBits,Origin,Size}` flat varyings from VS,
+ * the substituted `atlasSample(pageRef, formatBits, origin, size, uv)`
+ * call reads each header field inline from `headersU32` via
+ * `heap_drawIdx`. The wrapper-supplied `heap_drawIdx` u32 fn parameter
+ * carries the per-draw header index.
+ */
+function rewriteFsAtlasTexturesDirect(m, layout) {
+    if (layout.atlasTextureBindings.size === 0)
+        return m;
+    const used = new Set();
+    for (const v of m.values) {
+        if (v.kind !== "Entry" || v.entry.stage !== "fragment")
+            continue;
+        visitStmtExprs(v.entry.body, e => {
+            const hit = isAtlasSampleCall(e, layout.atlasTextureBindings);
+            if (hit !== null)
+                used.add(hit.name);
+        });
+    }
+    if (used.size === 0)
+        return m;
+    const fieldByAtlas = new Map();
+    for (const f of layout.drawHeaderFields) {
+        if (f.kind !== "texture-ref" || f.textureBindingName === undefined)
+            continue;
+        const sub = f.textureSub;
+        let inner = fieldByAtlas.get(f.textureBindingName);
+        if (inner === undefined) {
+            inner = new Map();
+            fieldByAtlas.set(f.textureBindingName, inner);
+        }
+        inner.set(sub, byteOffsetOf(f.byteOffset));
+    }
+    const drawIdxExpr = { kind: "Var", var: { name: "heap_drawIdx", type: Tu32, mutable: false } };
+    const stride = layout.strideU32;
+    const headerU32At = (offU32) => item(headersU32, add(mul(drawIdxExpr, constU32(stride), Tu32), constU32(offU32), Tu32), Tu32);
+    const headerF32At = (offU32) => ({
+        kind: "CallIntrinsic",
+        op: { name: "bitcast<f32>", returnTypeOf: () => Tf32, pure: true,
+            emit: { glsl: "intBitsToFloat", wgsl: "bitcast<f32>" } },
+        args: [headerU32At(offU32)],
+        type: Tf32,
+    });
+    return mapStageEntryBodies(m, "fragment", body => mapStmt(body, {
+        expr: e => mapExpr(e, sub => {
+            const hit = isAtlasSampleCall(sub, layout.atlasTextureBindings);
+            if (hit === null)
+                return sub;
+            const fields = fieldByAtlas.get(hit.name);
+            const pageRefOff = fields.get("pageRef").u32;
+            const formatBitsOff = fields.get("formatBits").u32;
+            const originOff = fields.get("origin").u32;
+            const sizeOff = fields.get("size").u32;
+            const args = [
+                headerU32At(pageRefOff),
+                headerU32At(formatBitsOff),
+                newVec([headerF32At(originOff), headerF32At(originOff + 1)], Tvec2),
+                newVec([headerF32At(sizeOff), headerF32At(sizeOff + 1)], Tvec2),
+                hit.uv,
+            ];
+            return { kind: "CallIntrinsic", op: ATLAS_SAMPLE_INTRINSIC, args, type: Tvec4 };
+        }),
+    }));
+}
+/**
+ * FS uniform threading: thread `_h_<name>Ref` (and `_iidx`, when
+ * needed) varyings from VS to FS, and rewrite FS `ReadInput("Uniform",
+ * X)` into a heap-load through the threaded ref.
+ *
+ * This pass stays at Module-level rather than going through the
+ * Effect.substitute API for two reasons:
+ *
+ *   1. The VS side adds *Stmts* (WriteOutput) and entry parameters,
+ *      not Expr substitutions — outside the substitute API's scope.
+ *   2. Even the FS Expr substitution can't run pre-merge: the
+ *      frontend hides `return { outColor: ... }` field exprs in a
+ *      `_record` carrier that mapExpr doesn't traverse, so any
+ *      uniform read inside the returned record is invisible until
+ *      `liftReturns` rewrites it into explicit WriteOutput stmts.
+ *      We run after liftReturns; the matching VS-side write injection
+ *      lives in the same pass to keep the `_h_<X>Ref` names paired.
+ */
+function rewriteFsUniforms(m, layout) {
+    // Scan FS bodies for uniform-scope reads.
+    const used = new Set();
+    for (const v of m.values) {
+        if (v.kind !== "Entry" || v.entry.stage !== "fragment")
+            continue;
+        for (const sn of readInputs(v.entry.body).values()) {
+            if (sn.scope === "Uniform")
+                used.add(sn.name);
+        }
+    }
+    if (used.size === 0)
+        return m;
+    const fieldByName = new Map(layout.drawHeaderFields.map(f => [f.name, f]));
+    const drawIdxExpr = drawIdxExprFor(layout);
+    const threadParams = [];
+    const vsWrites = [];
+    const fsSubst = new Map();
+    for (const name of used) {
+        const f = fieldByName.get(name);
+        if (f === undefined || f.kind !== "uniform-ref")
+            continue;
+        const refParamName = `_h_${name}Ref`;
+        const refExprWriter = loadHeaderRef(drawIdxExpr, f.byteOffset, layout.strideU32);
+        threadParams.push({
+            name: refParamName, type: Tu32, semantic: refParamName,
+            decorations: [{ kind: "Interpolation", mode: "flat" }],
+        });
+        vsWrites.push({
+            kind: "WriteOutput", name: refParamName,
+            value: { kind: "Expr", value: refExprWriter },
+        });
+        const refReadFs = readScope("Input", refParamName, Tu32);
+        const value = layout.perInstanceUniforms.has(name)
+            ? loadInstanceByRef(refReadFs, readScope("Input", "_iidx", Tu32), f.uniformWgslType ?? "")
+            : loadUniformByRef(refReadFs, f.uniformWgslType ?? "");
+        fsSubst.set(name, value);
+    }
+    // Per-instance: also thread iidx if any FS-used uniform is per-instance.
+    const fsHasPerInstance = [...used].some(n => layout.perInstanceUniforms.has(n));
+    if (fsHasPerInstance) {
+        threadParams.push({
+            name: "_iidx", type: Tu32, semantic: "_iidx",
+            decorations: [{ kind: "Interpolation", mode: "flat" }],
+        });
+        vsWrites.push({
+            kind: "WriteOutput", name: "_iidx",
+            value: { kind: "Expr", value: readScope("Builtin", "instance_index", Tu32) },
+        });
+    }
+    let out = m;
+    out = addInterstageParams(out, threadParams);
+    out = prependVsBodyStmts(out, vsWrites);
+    out = substituteInputsInStage(out, "fragment", "Uniform", n => fsSubst.get(n));
+    return out;
+}
+/**
+ * Per-instance uniform load: indexes into a packed array allocation
+ * by `iidx` (instance index in [0, instanceCount-1]).
+ */
+function loadInstanceByRef(refIdent, iidx, wgslType) {
+    const baseV4 = div(add(refIdent, constU32(16), Tu32), constU32(16), Tu32);
+    const baseF32 = div(add(refIdent, constU32(16), Tu32), constU32(4), Tu32);
+    switch (wgslType) {
+        case "mat4x4<f32>": {
+            const start = add(baseV4, mul(iidx, constU32(4), Tu32), Tu32);
+            const cols = [];
+            for (let i = 0; i < 4; i++) {
+                cols.push(item(heapV4f, i === 0 ? start : add(start, constU32(i), Tu32), Tvec4));
+            }
+            return matFromRows(cols, { kind: "Matrix", element: Tf32, rows: 4, cols: 4 });
+        }
+        case "vec4<f32>":
+            return item(heapV4f, add(baseV4, iidx, Tu32), Tvec4);
+        case "vec3<f32>": {
+            const off = mul(iidx, constU32(3), Tu32);
+            const base = add(baseF32, off, Tu32);
+            return newVec([
+                item(heapF32, base, Tf32),
+                item(heapF32, add(base, constU32(1), Tu32), Tf32),
+                item(heapF32, add(base, constU32(2), Tu32), Tf32),
+            ], Tvec3);
+        }
+        case "vec2<f32>": {
+            const off = mul(iidx, constU32(2), Tu32);
+            const base = add(baseF32, off, Tu32);
+            return newVec([
+                item(heapF32, base, Tf32),
+                item(heapF32, add(base, constU32(1), Tu32), Tf32),
+            ], Tvec2);
+        }
+        case "f32":
+            return item(heapF32, add(baseF32, iidx, Tu32), Tf32);
+        case "u32":
+            return item(heapU32, add(baseF32, iidx, Tu32), Tu32);
+        default:
+            throw new Error(`heapEffectIR: no IR loader for per-instance uniform type '${wgslType}'`);
+    }
+}
+// `wgslTypeFromHeapField` will be used once we propagate proper types
+// onto the synthesised threading params; for now we only emit u32-typed
+// refs, so the type lookup isn't called from the generated IR. Keep
+// the import alive.
+void wgslTypeFromHeapField;
+// ─── FS atlas-texture rewrite (textureSample → atlasSample) ─────────
+//
+// For each atlas-routed binding actually called from FS as
+// `textureSample(<name>, smp, uv)`:
+//   - Thread 4 flat-interpolated VsOut varyings: `_h_<name>PageRef`,
+//     `_h_<name>FormatBits`, `_h_<name>Origin`, `_h_<name>Size`.
+//   - VS body writes them from the drawHeader (`headersU32` reads,
+//     bitcast for the f32 fields since the headers buffer is u32).
+//   - FS substitutes the textureSample call with a `CallIntrinsic`
+//     emitting `atlasSample(pageRef, formatBits, origin, size, uv)`.
+//
+// The `atlasSample` helper is declared by the prelude string the post-
+// emit injector prepends to the FS source — the IR doesn't need to
+// know about its body.
+/** Synthetic `atlasSample` intrinsic for the IR's CallIntrinsic emit hook. */
+const ATLAS_SAMPLE_INTRINSIC = {
+    name: "atlasSample",
+    returnTypeOf: () => Tvec4,
+    pure: true,
+    emit: { glsl: "atlasSample", wgsl: "atlasSample" },
+};
+function isAtlasSampleCall(e, atlasNames) {
+    if (e.kind !== "CallIntrinsic")
+        return null;
+    if (e.op.emit.wgsl !== "textureSample")
+        return null;
+    if (e.args.length < 2)
+        return null;
+    const tex = e.args[0];
+    // Pre-legaliseTypes the sampler arg is a Var or ReadInput referencing
+    // the Sampler ValueDef by its schema name. Match either shape.
+    let name;
+    if (tex.kind === "ReadInput" && tex.scope === "Uniform")
+        name = tex.name;
+    else if (tex.kind === "Var")
+        name = tex.var.name;
+    if (name === undefined)
+        return null;
+    // The schema (extracted post-legaliseTypes) names the binding
+    // `<sampler>_view`; the IR pre-pass sees the user-written sampler
+    // name. Match either form so the rewrite fires regardless of which
+    // shape the layout's atlas-binding set carries.
+    const matched = atlasNames.has(name) ? name :
+        atlasNames.has(`${name}_view`) ? `${name}_view` :
+            undefined;
+    if (matched === undefined)
+        return null;
+    // texture(sampler, uv) parses as a 2-arg call (no separate sampler/
+    // texture pair until splitWgslSamplers runs). The uv is args[1].
+    const uv = e.args[1];
+    return { name: matched, uv };
+}
+function rewriteFsAtlasTextures(m, layout) {
+    if (layout.atlasTextureBindings.size === 0)
+        return m;
+    // Scan FS bodies to find which atlas bindings are actually used.
+    const used = new Set();
+    for (const v of m.values) {
+        if (v.kind !== "Entry" || v.entry.stage !== "fragment")
+            continue;
+        visitStmtExprs(v.entry.body, e => {
+            const hit = isAtlasSampleCall(e, layout.atlasTextureBindings);
+            if (hit !== null)
+                used.add(hit.name);
+        });
+    }
+    if (used.size === 0)
+        return m;
+    const fieldByAtlas = new Map();
+    for (const f of layout.drawHeaderFields) {
+        if (f.kind !== "texture-ref" || f.textureBindingName === undefined)
+            continue;
+        const sub = f.textureSub;
+        let inner = fieldByAtlas.get(f.textureBindingName);
+        if (inner === undefined) {
+            inner = new Map();
+            fieldByAtlas.set(f.textureBindingName, inner);
+        }
+        inner.set(sub, byteOffsetOf(f.byteOffset));
+    }
+    const drawIdxExpr = drawIdxExprFor(layout);
+    const threadParams = [];
+    const vsWrites = [];
+    for (const name of used) {
+        const v = atlasVaryingNames(name);
+        const fields = fieldByAtlas.get(name);
+        const stride = layout.strideU32;
+        const headerU32At = (offU32) => item(headersU32, add(mul(drawIdxExpr, constU32(stride), Tu32), constU32(offU32), Tu32), Tu32);
+        const headerF32At = (offU32) => ({
+            kind: "CallIntrinsic",
+            op: { name: "bitcast<f32>", returnTypeOf: () => Tf32, pure: true,
+                emit: { glsl: "intBitsToFloat", wgsl: "bitcast<f32>" } },
+            args: [headerU32At(offU32)],
+            type: Tf32,
+        });
+        const pageRefOff = fields.get("pageRef").u32;
+        const formatBitsOff = fields.get("formatBits").u32;
+        const originOff = fields.get("origin").u32;
+        const sizeOff = fields.get("size").u32;
+        threadParams.push({
+            name: v.pageRef, type: Tu32, semantic: v.pageRef,
+            decorations: [{ kind: "Interpolation", mode: "flat" }],
+        });
+        threadParams.push({
+            name: v.formatBits, type: Tu32, semantic: v.formatBits,
+            decorations: [{ kind: "Interpolation", mode: "flat" }],
+        });
+        threadParams.push({
+            name: v.origin, type: Tvec2, semantic: v.origin,
+            decorations: [{ kind: "Interpolation", mode: "flat" }],
+        });
+        threadParams.push({
+            name: v.size, type: Tvec2, semantic: v.size,
+            decorations: [{ kind: "Interpolation", mode: "flat" }],
+        });
+        vsWrites.push({ kind: "WriteOutput", name: v.pageRef,
+            value: { kind: "Expr", value: headerU32At(pageRefOff) } });
+        vsWrites.push({ kind: "WriteOutput", name: v.formatBits,
+            value: { kind: "Expr", value: headerU32At(formatBitsOff) } });
+        vsWrites.push({ kind: "WriteOutput", name: v.origin,
+            value: { kind: "Expr", value: newVec([headerF32At(originOff), headerF32At(originOff + 1)], Tvec2) } });
+        vsWrites.push({ kind: "WriteOutput", name: v.size,
+            value: { kind: "Expr", value: newVec([headerF32At(sizeOff), headerF32At(sizeOff + 1)], Tvec2) } });
+    }
+    // Rewrite FS textureSample(...) → atlasSample(...) for atlas-routed names.
+    let out = m;
+    out = addInterstageParams(out, threadParams);
+    out = prependVsBodyStmts(out, vsWrites);
+    out = mapStageEntryBodies(out, "fragment", body => mapStmt(body, {
+        expr: e => mapExpr(e, sub => {
+            const hit = isAtlasSampleCall(sub, layout.atlasTextureBindings);
+            if (hit === null)
+                return sub;
+            const v = atlasVaryingNames(hit.name);
+            const args = [
+                readScope("Input", v.pageRef, Tu32),
+                readScope("Input", v.formatBits, Tu32),
+                readScope("Input", v.origin, Tvec2),
+                readScope("Input", v.size, Tvec2),
+                hit.uv,
+            ];
+            return { kind: "CallIntrinsic", op: ATLAS_SAMPLE_INTRINSIC, args, type: Tvec4 };
+        }),
+    }));
+    return out;
+}
+/** Visit every Expr inside a Stmt tree, read-only. */
+function visitStmtExprs(s, fn) {
+    mapStmt(s, {
+        expr: e => mapExpr(e, sub => { fn(sub); return sub; }),
+    });
+}
+function byteOffsetOf(byteOffset) {
+    return { u32: byteOffset / 4 };
+}
+export function compileHeapEffectIR(userEffect, layout, compileOptions, mode = "standalone") {
+    // VS-side input/uniform substitution: pure read-rewrites, expressed
+    // via the Effect.substitute API.
+    const heapShaped = userEffect.substitute({
+        vertex: {
+            inputs: buildVertexAttributeMap(layout),
+            uniforms: buildVertexUniformMap(layout),
+        },
+    });
+    // Build a single Module from all of the effect's stages by
+    // composing them. Then apply the remaining heap rewrites to that
+    // combined module. The composed module preserves stage entry
+    // boundaries.
+    let combined = mergeStages(heapShaped);
+    // Lift `return { ... }` into explicit WriteOutput stmts so the FS
+    // uniform pass below can see (and rewrite) ReadInputs that feed the
+    // returned record fields. Without this, ObjectLiteral exprs hide
+    // their fields on a `_record` carrier that mapExpr doesn't traverse.
+    combined = liftReturns(combined);
+    combined = injectVsBuiltins(combined);
+    // FS uniform threading must run BEFORE atlas FS rewriting in source
+    // order? No — order historically: atlas first, then uniforms. Both
+    // add interstage params and VS WriteOutput stmts; their write sets
+    // are disjoint, so the order is interchangeable in practice. Keep
+    // the original order for byte-for-byte equivalence.
+    if (mode === "family-member") {
+        // Direct heap-read FS path: substitute uniforms / atlas samples
+        // directly using the wrapper-supplied `heap_drawIdx` / `instId`
+        // u32 fn parameters. No VS→FS varying threading on the heap side.
+        combined = rewriteFsAtlasTexturesDirect(combined, layout);
+        combined = rewriteFsUniformsDirect(combined, layout);
+    }
+    else {
+        combined = rewriteFsAtlasTextures(combined, layout);
+        combined = rewriteFsUniforms(combined, layout);
+    }
+    // Add heap storage buffers as bindings.
+    combined = { ...combined, values: [...heapStorageBufferDecls(), ...combined.values] };
+    // Skip frontend resolveHoles / liftReturns since they expect the
+    // template not to have been pre-modified. Compile via compileModule
+    // directly which still runs assignLocations + linkCrossStage + DCE +
+    // emit. (`makeStage` wrapper goes through Effect.compile which calls
+    // compileModule too.)
+    const compiled = compileModule(combined, compileOptions);
+    const vsStage = compiled.stages.find(s => s.stage === "vertex");
+    const fsStage = compiled.stages.find(s => s.stage === "fragment");
+    let vsSrc = vsStage?.source ?? "";
+    let fsSrc = fsStage?.source ?? "";
+    if (vsSrc.length > 0) {
+        vsSrc = mode === "family-member"
+            ? applyFamilyMemberShape(vsSrc)
+            : applyMegacallToEmittedVs(vsSrc);
+    }
+    if (mode === "family-member" && fsSrc.length > 0) {
+        fsSrc = applyFamilyMemberFsShape(fsSrc);
+    }
+    if (layout.atlasTextureBindings.size > 0) {
+        // Strip user-emitted texture/sampler binding decls for atlas-routed
+        // names — the IR's FS still emits `var <name>: texture_2d<f32>;` even
+        // though we've replaced every textureSample call referencing it.
+        // The render BGL has no slot for these bindings.
+        if (vsSrc.length > 0)
+            vsSrc = stripAtlasTextureSamplerDecls(vsSrc, layout.atlasTextureBindings);
+        if (fsSrc.length > 0) {
+            fsSrc = stripAtlasTextureSamplerDecls(fsSrc, layout.atlasTextureBindings);
+            fsSrc = prependAtlasPrelude(fsSrc);
+        }
+    }
+    return {
+        vs: vsSrc,
+        fs: fsSrc,
+        vsEntry: vsStage?.entryName ?? "vs",
+        fsEntry: fsStage?.entryName ?? "fs",
+        preludeWgsl: "",
+    };
+}
+/**
+ * Drop `@group(N) @binding(M) var <name>: texture_*<...>;` and the
+ * matching sampler decl(s) for atlas-routed binding names. We keep
+ * other texture/sampler decls intact since they belong to the standalone
+ * texture path.
+ *
+ * Sampler-name detection: WGSL emit pairs samplers with their textures
+ * positionally in the IR; the user's WGSL can name them anything. To
+ * be safe we also strip any sampler decl that is *only* referenced by
+ * a now-rewritten textureSample call. Rather than do data-flow, we
+ * pattern-match: if the FS has zero remaining textureSample calls AND
+ * a single sampler decl, drop it. Otherwise leave it. Practically the
+ * v1 surface has 0..1 atlas-routed binding per shader, so this is
+ * sufficient.
+ */
+function stripAtlasTextureSamplerDecls(src, atlasNames) {
+    let out = src;
+    for (const name of atlasNames) {
+        const re = new RegExp(`@group\\(\\d+\\)\\s*@binding\\(\\d+\\)\\s*var\\s+${name}\\s*:\\s*texture_\\w+(?:<[^>]*>)?\\s*;\\s*`, "g");
+        out = out.replace(re, "");
+    }
+    // If textureSample no longer appears anywhere, drop the (single)
+    // sampler decl too. Conservative: we leave samplers alone if they
+    // could still be in use.
+    if (!/textureSample\s*\(/.test(out)) {
+        out = out.replace(/@group\(\d+\)\s*@binding\(\d+\)\s*var\s+\w+\s*:\s*sampler(?:_comparison)?\s*;\s*/g, "");
+    }
+    return out;
+}
+/**
+ * Prepend the atlas texture-array decls + `atlasSample` helper to
+ * the emitted FS source. The helper signatures match the post-
+ * rewrite `atlasSample(pageRef, formatBits, origin, size, uv)`
+ * calls the IR pass produces.
+ */
+function prependAtlasPrelude(fs) {
+    // Delegate to the VS prelude generator — both stages need the same
+    // atlas-sample helper, and consolidating prevents the two copies
+    // from drifting.
+    return generateAtlasPrelude() + fs;
+}
+/**
+ * Megacall post-processing for the IR-emitted VS source. The IR uses
+ * `@builtin(vertex_index) vertex_index` + `@builtin(instance_index)
+ * instance_index` as the param names; rather than rewrite the IR
+ * (which would require building the binary-search loop in IR), we
+ * post-process the emitted WGSL: strip instance_index from the
+ * @vertex signature, rename vertex_index to emitIdx, and inject a
+ * prelude that defines `instance_index` + `vertex_index` as locals
+ * via the drawTable binary search so existing body references keep
+ * resolving. Also injects the drawTable/indexStorage binding decls.
+ */
+function applyMegacallToEmittedVs(vs) {
+    let s = vs;
+    // Megacall storage bindings only — the shared values (`heap_drawIdx`,
+    // `instId`, `vid`) are declared as `let` locals in the @vertex body
+    // and threaded as parameters into any helper fn that references them.
+    // No module-scope `var<private>` (Safari/WebKit rejects reading
+    // those from helper fn bodies).
+    const decl = `\n@group(0) @binding(4) var<storage, read> drawTable:       array<u32>;\n@group(0) @binding(5) var<storage, read> indexStorage:    array<u32>;\n@group(0) @binding(6) var<storage, read> firstDrawInTile: array<u32>;\n`;
+    // Locate the @vertex fn header, then balance parens manually since
+    // params can carry `@builtin(name)` decorations (regex `[^)]*` would
+    // halt at the first inner `)`).
+    const startRe = /@vertex\s+fn\s+(\w+)\s*\(/;
+    const startMatch = s.match(startRe);
+    if (startMatch === null)
+        return s;
+    const fnName = startMatch[1];
+    const paramOpen = startMatch.index + startMatch[0].length;
+    let depth = 1;
+    let i = paramOpen;
+    for (; i < s.length && depth > 0; i++) {
+        const c = s[i];
+        if (c === "(")
+            depth++;
+        else if (c === ")")
+            depth--;
+    }
+    if (depth !== 0)
+        return s;
+    const paramList = s.slice(paramOpen, i - 1);
+    // After the closing `)` we expect `\s*->\s*<retType>\s*\{`.
+    const tailRe = /\s*->\s*([\w<>,\s]+?)\s*\{/y;
+    tailRe.lastIndex = i;
+    const tailMatch = tailRe.exec(s);
+    if (tailMatch === null)
+        return s;
+    const retType = tailMatch[1].trim();
+    const headerEnd = tailRe.lastIndex;
+    const headerStart = startMatch.index;
+    const params = paramList.split(",").map(p => p.trim()).filter(p => p.length > 0);
+    const kept = [];
+    for (const p of params) {
+        if (/@builtin\(\s*instance_index\s*\)/.test(p))
+            continue;
+        if (/@builtin\(\s*vertex_index\s*\)/.test(p)) {
+            kept.push("@builtin(vertex_index) emitIdx: u32");
+            continue;
+        }
+        kept.push(p);
+    }
+    // megacallSearchPrelude exposes `heap_drawIdx`, `instId`, `vid`. The
+    // IR rewrites bind header-selector reads to `heap_drawIdx` directly
+    // (loadHeaderRef takes an Expr arg); the remaining bindings expose
+    // `instance_index` (= per-RO instance idx) and `vertex_index` for
+    // user code that reads them by builtin.
+    const newHeader = `@vertex fn ${fnName}(${kept.join(", ")}) -> ${retType} {\n${megacallSearchPrelude()}  let instance_index: u32 = instId;\n  let vertex_index: u32 = vid;\n`;
+    s = s.slice(0, headerStart) + newHeader + s.slice(headerEnd);
+    s = threadMegacallParamsThroughHelpers(s);
+    return decl + s;
+}
+/**
+ * Family-member shape: the wrapper performs the megacall binary search
+ * and invokes this VS as a regular function, passing `heap_drawIdx`,
+ * `instId`, and `vid` as plain `u32` parameters. We rewrite the emitted
+ * @vertex signature so its first three params are exactly those names
+ * (no `@builtin` decorations), and inject local aliases so any body
+ * references to `vertex_index` / `instance_index` (the names produced
+ * by IR `Builtin` reads + `injectVsBuiltins`) resolve to the new
+ * parameters.
+ *
+ * The megacall storage-buffer bindings (drawTable / indexStorage /
+ * firstDrawInTile) are NOT emitted — the wrapper module owns them.
+ *
+ * `threadMegacallParamsThroughHelpers` is run unchanged: helper fns
+ * extracted by composeStages may still reference `heap_drawIdx`,
+ * `instId`, `vid` after CSE; threading the params is a pure text walk
+ * and does not depend on the entry shape.
+ */
+function applyFamilyMemberShape(vs) {
+    let s = vs;
+    // Locate the @vertex fn header, then balance parens manually since
+    // params can carry `@builtin(name)` decorations.
+    const startRe = /@vertex\s+fn\s+(\w+)\s*\(/;
+    const startMatch = s.match(startRe);
+    if (startMatch === null)
+        return s;
+    const fnName = startMatch[1];
+    const paramOpen = startMatch.index + startMatch[0].length;
+    let depth = 1;
+    let i = paramOpen;
+    for (; i < s.length && depth > 0; i++) {
+        const c = s[i];
+        if (c === "(")
+            depth++;
+        else if (c === ")")
+            depth--;
+    }
+    if (depth !== 0)
+        return s;
+    const paramList = s.slice(paramOpen, i - 1);
+    const tailRe = /\s*->\s*([\w<>,\s]+?)\s*\{/y;
+    tailRe.lastIndex = i;
+    const tailMatch = tailRe.exec(s);
+    if (tailMatch === null)
+        return s;
+    const retType = tailMatch[1].trim();
+    const headerEnd = tailRe.lastIndex;
+    const headerStart = startMatch.index;
+    // Drop ALL @vertex fn params: the megacall builtins
+    // (instance_index, vertex_index) are not needed because we re-bind
+    // them from incoming u32 args below, and any @location(...) attribute
+    // params have already had their body reads substituted with
+    // heap-load expressions by `substituteInputsInStage` — the params
+    // themselves are dead. The wrapper calls the helper with exactly
+    // three megacall args.
+    const newParamList = ["heap_drawIdx: u32", "instId: u32", "vid: u32"].join(", ");
+    void paramList;
+    // Body aliases: IR-emitted body code reads `vertex_index` and
+    // `instance_index` (the names of the @builtin params standalone mode
+    // adds). Re-bind those names to the incoming `vid` / `instId`.
+    const bodyAliases = `  let instance_index: u32 = instId;\n` +
+        `  let vertex_index: u32 = vid;\n`;
+    const newHeader = `@vertex fn ${fnName}(${newParamList}) -> ${retType} {\n${bodyAliases}`;
+    s = s.slice(0, headerStart) + newHeader + s.slice(headerEnd);
+    s = threadMegacallParamsThroughHelpers(s);
+    return s;
+}
+/**
+ * Family-member FS shape: rewrite the emitted `@fragment fn` signature
+ * to take `heap_drawIdx: u32, instId: u32` as additional plain u32 fn
+ * parameters (after the existing FsIn struct param). The body uses
+ * `heap_drawIdx` / `instId` as `Var` reads — the wrapper passes them
+ * in from flat-interpolated u32 inputs on the family FsIn.
+ *
+ * `threadMegacallParamsThroughHelpers` is run as well in case any FS
+ * helper fn extracted by composeStages references those identifiers.
+ */
+function applyFamilyMemberFsShape(fs) {
+    let s = fs;
+    const startRe = /@fragment\s+fn\s+(\w+)\s*\(/;
+    const startMatch = s.match(startRe);
+    if (startMatch === null)
+        return s;
+    const fnName = startMatch[1];
+    const paramOpen = startMatch.index + startMatch[0].length;
+    let depth = 1;
+    let i = paramOpen;
+    for (; i < s.length && depth > 0; i++) {
+        const c = s[i];
+        if (c === "(")
+            depth++;
+        else if (c === ")")
+            depth--;
+    }
+    if (depth !== 0)
+        return s;
+    const paramList = s.slice(paramOpen, i - 1);
+    const tailRe = /\s*->\s*([\w<>,\s]+?)\s*\{/y;
+    tailRe.lastIndex = i;
+    const tailMatch = tailRe.exec(s);
+    if (tailMatch === null)
+        return s;
+    const retType = tailMatch[1].trim();
+    const headerEnd = tailRe.lastIndex;
+    const headerStart = startMatch.index;
+    const existing = paramList.split(",").map(p => p.trim()).filter(p => p.length > 0);
+    const newParamList = [...existing, "heap_drawIdx: u32", "instId: u32"].join(", ");
+    const newHeader = `@fragment fn ${fnName}(${newParamList}) -> ${retType} {\n`;
+    s = s.slice(0, headerStart) + newHeader + s.slice(headerEnd);
+    s = threadMegacallParamsThroughHelpers(s);
+    return s;
+}
+/**
+ * Post-emit text rewrite that threads `heap_drawIdx`, `instId`, and
+ * `vid` through any helper function whose body references them. The
+ * IR's composeStages pass extracts same-stage entries into helper fns
+ * (`fn _<name>(s_in: <State>) -> <State>`) called from the wrapper
+ * `@vertex` fn. Those helpers can carry CSE-extracted expressions that
+ * reference `heap_drawIdx`/`instId`/`vid` — declared as locals in the
+ * wrapper, those identifiers are out of scope inside helpers, so we
+ * append them as `u32` parameters and pass them at every call site.
+ *
+ * Module-scope `var<private>` would also work but Safari/WebKit's WGSL
+ * parser rejects helper-fn reads of module-scope private vars in some
+ * configurations; explicit parameter passing is portable across all
+ * conforming WGSL implementations.
+ */
+function threadMegacallParamsThroughHelpers(src) {
+    const idents = ["heap_drawIdx", "instId", "vid"];
+    // Find every `fn _<name>(...) -> <ret> { ... }` declaration. We only
+    // touch fns whose name starts with `_` (the convention used by
+    // extractFusedEntry's helpers); user code with a leading-underscore
+    // name is not expected in IR-emitted WGSL.
+    const fnRe = /\bfn\s+(_\w+)\s*\(/g;
+    const edits = [];
+    let m;
+    while ((m = fnRe.exec(src)) !== null) {
+        const name = m[1];
+        const paramOpen = m.index + m[0].length; // position after `(`
+        // Balance parens to find the matching `)`.
+        let depth = 1;
+        let i = paramOpen;
+        for (; i < src.length && depth > 0; i++) {
+            const c = src[i];
+            if (c === "(")
+                depth++;
+            else if (c === ")")
+                depth--;
+        }
+        if (depth !== 0)
+            continue;
+        const paramClose = i - 1; // position of the matching `)`
+        // Find the function body `{ ... }` after the return type.
+        const afterParens = src.slice(i);
+        const braceIdx = afterParens.indexOf("{");
+        if (braceIdx < 0)
+            continue;
+        const bodyOpen = i + braceIdx;
+        let bdepth = 1;
+        let j = bodyOpen + 1;
+        for (; j < src.length && bdepth > 0; j++) {
+            const c = src[j];
+            if (c === "{")
+                bdepth++;
+            else if (c === "}")
+                bdepth--;
+        }
+        if (bdepth !== 0)
+            continue;
+        const bodyClose = j - 1;
+        const body = src.slice(bodyOpen + 1, bodyClose);
+        const needed = idents.filter(id => new RegExp(`\\b${id}\\b`).test(body));
+        if (needed.length === 0)
+            continue;
+        edits.push({ name, needed, paramOpen, paramClose });
+    }
+    if (edits.length === 0)
+        return src;
+    // Apply param-list edits from the back so earlier offsets stay valid.
+    const helperNeeds = new Map();
+    let out = src;
+    for (let k = edits.length - 1; k >= 0; k--) {
+        const e = edits[k];
+        helperNeeds.set(e.name, e.needed);
+        const existing = out.slice(e.paramOpen, e.paramClose).trim();
+        const extra = e.needed.map(id => `${id}: u32`).join(", ");
+        const newParams = existing.length === 0 ? extra : `${existing}, ${extra}`;
+        out = out.slice(0, e.paramOpen) + newParams + out.slice(e.paramClose);
+    }
+    // Rewrite call sites for each touched helper. Match `<name>(<args>)`
+    // and append the corresponding identifiers. Match the bare name so
+    // we don't accidentally hit a substring inside another identifier.
+    for (const [name, needed] of helperNeeds) {
+        const callRe = new RegExp(`\\b${name}\\s*\\(`, "g");
+        let result = "";
+        let lastIdx = 0;
+        let cm;
+        while ((cm = callRe.exec(out)) !== null) {
+            const callOpen = cm.index + cm[0].length;
+            // Skip if this is the `fn <name>(` declaration itself (preceded
+            // by `fn`).
+            const before = out.slice(0, cm.index).trimEnd();
+            if (/\bfn$/.test(before))
+                continue;
+            // Balance parens to find matching `)`.
+            let depth = 1;
+            let p = callOpen;
+            for (; p < out.length && depth > 0; p++) {
+                const c = out[p];
+                if (c === "(")
+                    depth++;
+                else if (c === ")")
+                    depth--;
+            }
+            if (depth !== 0)
+                continue;
+            const callClose = p - 1;
+            const args = out.slice(callOpen, callClose).trim();
+            const extras = needed.join(", ");
+            const newArgs = args.length === 0 ? extras : `${args}, ${extras}`;
+            result += out.slice(lastIdx, callOpen) + newArgs + ")";
+            lastIdx = p; // p is one past the `)`
+        }
+        if (lastIdx > 0) {
+            result += out.slice(lastIdx);
+            out = result;
+        }
+    }
+    return out;
+}
+function mergeStages(eff) {
+    // Combine all stages' templates into one module by concatenating
+    // their values + types arrays. The compile pipeline's composeStages
+    // pass will fuse same-stage entries afterwards.
+    const types = eff.stages.flatMap(s => [...s.template.types]);
+    const values = eff.stages.flatMap(s => [...s.template.values]);
+    return { types, values };
+}
+//# sourceMappingURL=heapEffectIR.js.map