@aardworx/wombat.rendering 0.9.13 → 0.9.15

package/src/runtime/derivedModes/modeKeyCpu.ts

@@ -0,0 +1,213 @@
+// Per-RO modeKey production (CPU side).
+//
+// Task 1 of the derived-modes work: turn a `PipelineState` (the
+// aval-shaped struct from core/pipelineState.ts) plus the surrounding
+// `FramebufferSignature` into a canonical `PipelineStateDescriptor`
+// and its bitfield-encoded modeKey. Subscribes to every input aval
+// so that mutating a `cval` for cullMode / blendFactor / etc. fires
+// a dirty callback, the bucket flushes the changed RO's modeKey to
+// GPU on the next frame, and the partition kernel re-buckets the
+// record into its new slot. No rule-IR machinery yet — Task 2 layers
+// `derivedMode(...)` on top of this.
+
+import type {
+  DepthState,
+  PipelineState,
+  RasterizerState,
+  StencilState,
+  BlendState as PsBlendState,
+  BlendComponentState as PsBlendComponentState,
+} from "../../core/pipelineState.js";
+import type { FramebufferSignature } from "../../core/framebufferSignature.js";
+import type { aval, IDisposable } from "@aardworx/wombat.adaptive";
+import { addMarkingCallback } from "@aardworx/wombat.adaptive";
+
+import {
+  encodeModeKey,
+  DEFAULT_DESCRIPTOR,
+  DEFAULT_ATTACHMENT_BLEND,
+  type PipelineStateDescriptor,
+  type AttachmentBlend,
+  type BlendComponent,
+  type DepthSlice,
+} from "../pipelineCache/index.js";
+
+/**
+ * One-shot snapshot of the current values of every leaf aval in a
+ * `PipelineState`, mapped onto a `PipelineStateDescriptor` whose
+ * `attachments` array follows the framebuffer's `colorNames` order
+ * (missing entries default to non-blended / 0xF writeMask, matching
+ * `BLEND_DEFAULT` semantics in heapScene).
+ */
+export function snapshotDescriptor(
+  ps: PipelineState | undefined,
+  signature: FramebufferSignature,
+): PipelineStateDescriptor {
+  if (ps === undefined) {
+    return buildAttachmentsFor(DEFAULT_DESCRIPTOR, signature);
+  }
+  const r = ps.rasterizer;
+  const depth = ps.depth !== undefined ? snapshotDepth(ps.depth, signature) : undefined;
+  const atc = ps.alphaToCoverage !== undefined ? ps.alphaToCoverage.force(/* allow-force */) : false;
+  const blends = ps.blends !== undefined ? ps.blends.force(/* allow-force */) : undefined;
+  const attachments: AttachmentBlend[] = signature.colorNames.map((name) => {
+    const b = blends?.tryFind(name);
+    return b !== undefined ? snapshotAttachment(b) : DEFAULT_ATTACHMENT_BLEND;
+  });
+  const out: PipelineStateDescriptor = {
+    topology:         r.topology.force(/* allow-force */),
+    stripIndexFormat: stripFormatFor(r.topology.force(/* allow-force */)),
+    frontFace:        r.frontFace.force(/* allow-force */),
+    cullMode:         r.cullMode.force(/* allow-force */),
+    ...(depth !== undefined ? { depth } : {}),
+    attachments,
+    alphaToCoverage: atc,
+  };
+  return out;
+}
+
+function buildAttachmentsFor(
+  base: PipelineStateDescriptor,
+  signature: FramebufferSignature,
+): PipelineStateDescriptor {
+  // When `ps` is undefined, every named attachment uses defaults.
+  return {
+    ...base,
+    attachments: signature.colorNames.map(() => DEFAULT_ATTACHMENT_BLEND),
+  };
+}
+
+function snapshotDepth(d: DepthState, signature: FramebufferSignature): DepthSlice | undefined {
+  if (signature.depthStencil === undefined) return undefined;
+  return {
+    write:   d.write.force(/* allow-force */),
+    compare: d.compare.force(/* allow-force */),
+    clamp:   d.clamp !== undefined ? d.clamp.force(/* allow-force */) : false,
+  };
+}
+
+function snapshotAttachment(b: PsBlendState): AttachmentBlend {
+  const color = snapshotBlendComponent(b.color);
+  const alpha = snapshotBlendComponent(b.alpha);
+  const writeMask = b.writeMask.force(/* allow-force */) & 0xF;
+  // "Blend enabled" derives from whether the components differ from
+  // the no-op (src=one, dst=zero, op=add) shape. This mirrors WebGPU's
+  // contract: a blend descriptor MUST be present for blending to occur,
+  // and the no-op shape is the canonical disabled state.
+  const enabled =
+    !(color.srcFactor === "one" && color.dstFactor === "zero" && color.operation === "add" &&
+      alpha.srcFactor === "one" && alpha.dstFactor === "zero" && alpha.operation === "add");
+  return { enabled, color, alpha, writeMask };
+}
+
+function snapshotBlendComponent(c: PsBlendComponentState): BlendComponent {
+  return {
+    srcFactor: c.srcFactor.force(/* allow-force */),
+    dstFactor: c.dstFactor.force(/* allow-force */),
+    operation: c.operation.force(/* allow-force */),
+  };
+}
+
+function stripFormatFor(topology: GPUPrimitiveTopology): GPUIndexFormat | undefined {
+  return topology === "line-strip" || topology === "triangle-strip" ? "uint32" : undefined;
+}
+
+// ─── Reactive tracker ──────────────────────────────────────────────────
+
+/**
+ * Subscribes to every leaf aval in a `PipelineState` and invokes
+ * `onDirty` when any of them marks. Use one tracker per RO; call
+ * `recompute()` after a dirty signal to refresh the descriptor +
+ * modeKey, and `dispose()` when the RO is removed.
+ *
+ * Reactive cvals are the typical case; constant avals also "subscribe"
+ * — `addMarkingCallback` is a no-op for an aval that never marks, so
+ * no special-casing required.
+ */
+export class ModeKeyTracker implements IDisposable {
+  readonly ps: PipelineState | undefined;
+  readonly signature: FramebufferSignature;
+  private readonly onDirty: () => void;
+  private readonly subs: IDisposable[] = [];
+  private cachedDescriptor: PipelineStateDescriptor;
+  private cachedModeKey: bigint;
+
+  constructor(
+    ps: PipelineState | undefined,
+    signature: FramebufferSignature,
+    onDirty: () => void,
+  ) {
+    this.ps        = ps;
+    this.signature = signature;
+    this.onDirty   = onDirty;
+    this.cachedDescriptor = snapshotDescriptor(ps, signature);
+    this.cachedModeKey    = encodeModeKey(this.cachedDescriptor);
+    this.subscribeAll();
+  }
+
+  get descriptor(): PipelineStateDescriptor { return this.cachedDescriptor; }
+  get modeKey(): bigint { return this.cachedModeKey; }
+
+  /**
+   * Snapshot current aval values and rebuild descriptor + modeKey.
+   * Returns true iff the modeKey changed (so the bucket only needs to
+   * upload + repartition when something actually moved).
+   */
+  recompute(): boolean {
+    const next = snapshotDescriptor(this.ps, this.signature);
+    const nextKey = encodeModeKey(next);
+    if (nextKey === this.cachedModeKey) return false;
+    this.cachedDescriptor = next;
+    this.cachedModeKey    = nextKey;
+    return true;
+  }
+
+  dispose(): void {
+    for (const s of this.subs) s.dispose();
+    this.subs.length = 0;
+  }
+
+  private subscribeAll(): void {
+    if (this.ps === undefined) return;
+    const ps = this.ps;
+    this.sub(ps.rasterizer.topology);
+    this.sub(ps.rasterizer.cullMode);
+    this.sub(ps.rasterizer.frontFace);
+    if (ps.rasterizer.depthBias !== undefined) this.sub(ps.rasterizer.depthBias);
+    if (ps.depth !== undefined) {
+      this.sub(ps.depth.write);
+      this.sub(ps.depth.compare);
+      if (ps.depth.clamp !== undefined) this.sub(ps.depth.clamp);
+    }
+    if (ps.stencil !== undefined) {
+      // Stencil enabled/reference/masks contribute to neither modeKey
+      // nor cache key (v1 carves stencil out — only the "enabled"
+      // flag matters and we surface it via snapshotDescriptor). We
+      // still subscribe so callers can repackage if a future change
+      // flips the enable bit.
+      this.subStencilFaces(ps.stencil);
+    }
+    if (ps.blends !== undefined) {
+      this.sub(ps.blends);
+      const map = ps.blends.force(/* allow-force */);
+      for (const [, bs] of map) this.subBlendState(bs);
+    }
+    if (ps.alphaToCoverage !== undefined) this.sub(ps.alphaToCoverage);
+    // blendConstant is dynamic state; not in the modeKey, no subscription needed.
+  }
+
+  private sub(a: aval<unknown>): void {
+    this.subs.push(addMarkingCallback(a, this.onDirty));
+  }
+
+  private subBlendState(b: PsBlendState): void {
+    this.sub(b.color.srcFactor); this.sub(b.color.dstFactor); this.sub(b.color.operation);
+    this.sub(b.alpha.srcFactor); this.sub(b.alpha.dstFactor); this.sub(b.alpha.operation);
+    this.sub(b.writeMask);
+  }
+
+  private subStencilFaces(_s: StencilState): void {
+    // v1: stencil isn't in the modeKey; deliberately not subscribing
+    // to face avals here. Wire up in v2 when stencil enters the key.
+  }
+}

package/src/runtime/derivedModes/partition.ts

@@ -0,0 +1,206 @@
+// Per-frame partition kernel — turns per-RO slot assignments into
+// per-slot draw metadata (record lists + cumulative emit prefix sums
+// + total indexCount × instanceCount per slot).
+//
+// Each bucket runs this kernel once per frame (dirty-gated). Output
+// feeds one `drawIndirect` per slot in the encode loop. Records routed
+// to the same slot share an indirect call; records in different slots
+// land in separate `setPipeline + drawIndirect` calls.
+//
+// The CPU function `partitionCPU` is the algorithmic reference and
+// the basis for unit tests. The WGSL kernel below mirrors the same
+// algorithm using two compute passes (histogram + scatter), with an
+// exclusive prefix scan on the histogram in between.
+//
+// Layout choices:
+//
+//   - Records are NOT physically re-sorted. Instead, the kernel writes
+//     a permuted index list `slotRecords[]` (one u32 per record), plus
+//     per-slot offsets `slotOffsets[]` and counts `slotCounts[]`. The
+//     existing drawTable stays put — the VS prelude does one extra
+//     read of `slotRecords[…]` to map from "drawIndex within slot"
+//     back to "drawIndex in the full drawTable."
+//   - `slotEmitPrefix[]` is the slot-local prefix-sum of
+//     `record.indexCount * record.instanceCount`. The VS prelude
+//     binary-searches THIS (per-slot) instead of the global table.
+//   - `slotTotalEmit[]` is the dispatched vertex count for that
+//     slot's drawIndirect.
+//
+// Buffers grow lazily; the partition wrapper resizes as slot counts
+// or record counts change. v1 caps slots per bucket at MAX_SLOTS
+// (256), exceeded slots throw with a friendly diagnostic.
+
+/** v1 cap on distinct pipeline slots per bucket. */
+export const MAX_SLOTS_PER_BUCKET = 256;
+
+export interface PartitionInput {
+  /** Per-record slot index (length = numRecords). */
+  readonly slot:          ReadonlyArray<number>;
+  /** Per-record (indexCount * instanceCount) — emit count contributed
+   * to its slot's drawIndirect vertex total. */
+  readonly emitCount:     ReadonlyArray<number>;
+  /** Number of distinct slots used. Slots outside this range produce
+   * an error. */
+  readonly slotCount:     number;
+}
+
+export interface PartitionResult {
+  /** Per-slot record count. Length = slotCount. */
+  readonly slotCounts:     Uint32Array;
+  /** Per-slot cumulative record-count prefix sum (length = slotCount + 1). */
+  readonly slotOffsets:    Uint32Array;
+  /** Per-slot total emit count (vertex count for the drawIndirect). */
+  readonly slotTotalEmit:  Uint32Array;
+  /** Permuted record indices, grouped by slot. Length = numRecords. */
+  readonly slotRecords:    Uint32Array;
+  /** Slot-local prefix sum of emitCounts, length = numRecords + slotCount.
+   * For slot s, the prefix array is
+   *   slotEmitPrefix[slotOffsets[s] + s ..= slotOffsets[s+1] + s]
+   * (one extra entry per slot for the upper bound). */
+  readonly slotEmitPrefix: Uint32Array;
+}
+
+/**
+ * Pure-CPU partition. Reference implementation for the WGSL kernel
+ * below. Order within a slot is stable wrt. input record order so
+ * tests can predict the permuted layout deterministically.
+ */
+export function partitionCPU(input: PartitionInput): PartitionResult {
+  const { slot, emitCount, slotCount } = input;
+  const numRecords = slot.length;
+  if (emitCount.length !== numRecords) {
+    throw new Error("partition: slot.length != emitCount.length");
+  }
+  if (slotCount > MAX_SLOTS_PER_BUCKET) {
+    throw new Error(`partition: slotCount ${slotCount} exceeds cap ${MAX_SLOTS_PER_BUCKET}`);
+  }
+
+  const slotCounts    = new Uint32Array(slotCount);
+  const slotOffsets   = new Uint32Array(slotCount + 1);
+  const slotTotalEmit = new Uint32Array(slotCount);
+  const slotRecords   = new Uint32Array(numRecords);
+  const slotEmitPrefix = new Uint32Array(numRecords + slotCount);
+
+  // Pass 1: histogram + per-slot emit sum.
+  for (let i = 0; i < numRecords; i++) {
+    const s = slot[i]!;
+    if (s < 0 || s >= slotCount) {
+      throw new Error(`partition: record ${i} has out-of-range slot ${s} (slotCount=${slotCount})`);
+    }
+    slotCounts[s] = slotCounts[s]! + 1;
+    slotTotalEmit[s] = slotTotalEmit[s]! + emitCount[i]!;
+  }
+
+  // Pass 2: exclusive prefix sum -> per-slot start offsets.
+  let acc = 0;
+  for (let s = 0; s < slotCount; s++) {
+    slotOffsets[s] = acc;
+    acc += slotCounts[s]!;
+  }
+  slotOffsets[slotCount] = acc;
+
+  // Pass 3: scatter — second pass over records to place each at its
+  // slot's current cursor, emitting the within-slot prefix sums.
+  const cursors = new Uint32Array(slotCount);
+  for (let i = 0; i < numRecords; i++) {
+    const s = slot[i]!;
+    const base    = slotOffsets[s]!;     // first slotRecords[] index for this slot
+    const prefBase = base + s;           // first slotEmitPrefix[] index for this slot
+    const k       = cursors[s]!;         // within-slot position
+    slotRecords[base + k] = i;
+    // Slot prefix sums are length (slotCounts[s] + 1): first entry
+    // is 0, then cumulative; the last entry equals slotTotalEmit[s].
+    if (k === 0) {
+      slotEmitPrefix[prefBase] = 0;
+    }
+    slotEmitPrefix[prefBase + k + 1] = slotEmitPrefix[prefBase + k]! + emitCount[i]!;
+    cursors[s] = k + 1;
+  }
+
+  return { slotCounts, slotOffsets, slotTotalEmit, slotRecords, slotEmitPrefix };
+}
+
+// ─── WGSL kernel ───────────────────────────────────────────────────────
+//
+// Mirrors `partitionCPU` in two compute passes:
+//
+//   Pass A (histogram):
+//     - one thread per record
+//     - read slot[i], emitCount[i]
+//     - atomicAdd(slotCounts[s], 1)
+//     - atomicAdd(slotTotalEmit[s], emitCount[i])
+//
+//   Pass B (scatter):
+//     - one thread per record
+//     - look up slot[i]
+//     - atomicAdd(cursors[s], 1) -> within-slot position k
+//     - slotRecords[slotOffsets[s] + k] = i
+//     - slotEmitPrefix[slotOffsets[s] + s + k + 1] is computed by a
+//       second scan pass (the within-slot prefix sum). v1 emits a
+//       serial CPU-side scan for simplicity; v2 should fold this into
+//       a per-slot warp-scan once buckets get large.
+//
+// Between the two passes, the CPU runs an exclusive scan on
+// `slotCounts` to produce `slotOffsets` (small array, ≤ MAX_SLOTS).
+//
+// The within-slot prefix scan in pass B uses a non-atomic read of
+// `slotEmitPrefix[prefBase + k]` — this is safe because the only
+// thread that writes index `prefBase + k` is the one whose `cursors`
+// fetchAdd returned `k`, and that thread completed before any thread
+// whose fetchAdd returned `k+1` could read `prefBase + k`.
+//
+// (Wait — that's only true if we have a memory barrier across
+// invocations, which WGSL doesn't guarantee inside one dispatch. In
+// practice v1 punts: the kernel emits the histogram + slotRecords,
+// and the host or a follow-up pass computes slotEmitPrefix. The
+// kernel below reflects this: it produces slotCounts, slotOffsets
+// (after host scan), slotTotalEmit, and slotRecords. The
+// slotEmitPrefix scan is performed in a third compute pass on the
+// already-scattered records.)
+
+export const PARTITION_HISTOGRAM_WGSL = /* wgsl */ `
+struct PartitionInputs {
+  numRecords:   u32,
+  slotCount:    u32,
+  _pad0:        u32,
+  _pad1:        u32,
+};
+@group(0) @binding(0) var<uniform>            uIn:           PartitionInputs;
+@group(0) @binding(1) var<storage, read>      slot:          array<u32>;
+@group(0) @binding(2) var<storage, read>      emitCount:     array<u32>;
+@group(0) @binding(3) var<storage, read_write> slotCounts:    array<atomic<u32>>;
+@group(0) @binding(4) var<storage, read_write> slotTotalEmit: array<atomic<u32>>;
+
+@compute @workgroup_size(64)
+fn histogram(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i >= uIn.numRecords) { return; }
+  let s = slot[i];
+  atomicAdd(&slotCounts[s], 1u);
+  atomicAdd(&slotTotalEmit[s], emitCount[i]);
+}
+`;
+
+export const PARTITION_SCATTER_WGSL = /* wgsl */ `
+struct PartitionInputs {
+  numRecords:   u32,
+  slotCount:    u32,
+  _pad0:        u32,
+  _pad1:        u32,
+};
+@group(0) @binding(0) var<uniform>            uIn:         PartitionInputs;
+@group(0) @binding(1) var<storage, read>      slot:        array<u32>;
+@group(0) @binding(2) var<storage, read>      slotOffsets: array<u32>;
+@group(0) @binding(3) var<storage, read_write> cursors:     array<atomic<u32>>;
+@group(0) @binding(4) var<storage, read_write> slotRecords: array<u32>;
+
+@compute @workgroup_size(64)
+fn scatter(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i >= uIn.numRecords) { return; }
+  let s    = slot[i];
+  let base = slotOffsets[s];
+  let k    = atomicAdd(&cursors[s], 1u);
+  slotRecords[base + k] = i;
+}
+`;

package/src/runtime/derivedModes/slotTable.ts

@@ -0,0 +1,153 @@
+// SlotTable — per-bucket (or per-scene-scope) registry that maps a
+// per-RO modeKey to a fixed slot index. Each slot owns one
+// `GPURenderPipeline` and one offset into a shared indirect-args
+// buffer. The partition kernel reads `modeKeyToSlot` to atomic-add a
+// record's counts into the right slot; encode iterates the slots and
+// emits one `drawIndirect` per slot (zero-record slots draw zero).
+//
+// Lifecycle:
+//
+//   - addRO(modeKey, descriptor) → ensures the slot exists.
+//     - On miss, allocates a new slot, fires `precompile` for the
+//       missing descriptor (or `getOrCreateSync` if `sync=true`).
+//   - removeRO(modeKey) → drops a refcount on the slot. v1 keeps the
+//     slot around even at refcount=0 (pipelines are cheap to retain
+//     and the future case is "RO reappears next frame").
+//   - ready(): Promise<void> → resolves when every registered slot
+//     has a linked pipeline.
+//   - syncLookup(modeKey) → slot index or -1 (used by partition kernel
+//     setup; missing slot is a logical error post-ready()).
+//
+// This module does NOT touch heapScene. It's a generic data structure
+// that the heap integration drives. Tests use a mock builder.
+
+import {
+  PipelineCache,
+  encodeModeKey as encodeMK,
+  type PipelineBuilder,
+  type PipelineStateDescriptor,
+} from "../pipelineCache/index.js";
+
+export interface SlotTableEntry {
+  readonly slotIndex: number;
+  readonly modeKey:   bigint;
+  readonly descriptor: PipelineStateDescriptor;
+  /** Refcount of ROs currently bound to this slot. v1 doesn't GC at 0. */
+  refCount: number;
+  /** Populated lazily by precompile; null until linked. */
+  pipeline: GPURenderPipeline | null;
+}
+
+/**
+ * One SlotTable per `(effect, textureSet)` bucket. The cache may be
+ * shared across buckets — pipelines that happen to share a descriptor
+ * (rare across distinct effects, since pipeline objects bake the
+ * shader / layout) still dedupe at the cache level if a builder
+ * legitimately produces the same GPURenderPipeline for the same
+ * descriptor key.
+ */
+export class SlotTable {
+  private readonly entries: SlotTableEntry[] = [];
+  private readonly byKey   = new Map<bigint, SlotTableEntry>();
+  /** Tracks pending precompile promises so ready() waits for the union. */
+  private pending: Set<Promise<unknown>> = new Set();
+  private version = 0;
+
+  constructor(
+    readonly cache: PipelineCache,
+    readonly builder: PipelineBuilder,
+  ) {}
+
+  /** Total slot count (linked or pending). */
+  get slotCount(): number { return this.entries.length; }
+  /**
+   * Bumps whenever a new slot is appended. Consumers (partition setup,
+   * indirect-args resize, GPU lookup-table upload) compare against
+   * their last-seen value to know when to re-upload.
+   */
+  get layoutVersion(): number { return this.version; }
+  /** Iterate all slots in slot-index order. */
+  *all(): IterableIterator<SlotTableEntry> { yield* this.entries; }
+
+  /**
+   * Add (or reference-bump) the slot for `descriptor`. Returns the
+   * existing slot if it already exists, otherwise creates one and
+   * kicks off compilation via `builder` (async by default).
+   *
+   * Pass `sync: true` on the runtime-mutation path to use
+   * `builder.createSync` instead — the returned slot's `pipeline` will
+   * be populated immediately.
+   */
+  addRO(
+    descriptor: PipelineStateDescriptor,
+    options: { sync?: boolean } = {},
+  ): SlotTableEntry {
+    const modeKey = encodeMK(descriptor);
+    const existing = this.byKey.get(modeKey);
+    if (existing !== undefined) {
+      existing.refCount += 1;
+      return existing;
+    }
+    const slotIndex = this.entries.length;
+    const entry: SlotTableEntry = {
+      slotIndex, modeKey, descriptor, refCount: 1, pipeline: null,
+    };
+    this.entries.push(entry);
+    this.byKey.set(modeKey, entry);
+    this.version += 1;
+    if (options.sync === true) {
+      entry.pipeline = this.cache.getOrCreateSync(this.builder, descriptor);
+    } else {
+      // Use the cache's precompile path so concurrent requests
+      // deduplicate. Tracks promise so `ready()` blocks until done.
+      const p = this.cache.precompile(this.builder, [descriptor]).then(() => {
+        const got = this.cache.lookup(modeKey);
+        if (got !== null) entry.pipeline = got;
+      });
+      this.pending.add(p);
+      // Auto-clean once settled.
+      p.finally(() => this.pending.delete(p));
+    }
+    return entry;
+  }
+
+  /** Drop one reference from the slot. Pipelines stay resident. */
+  removeRO(modeKey: bigint): void {
+    const e = this.byKey.get(modeKey);
+    if (e === undefined) return;
+    if (e.refCount > 0) e.refCount -= 1;
+  }
+
+  /**
+   * Move an RO from one slot to another (reactive aval mutation
+   * caused its modeKey to change). Returns the new (possibly newly
+   * allocated) slot.
+   */
+  rebindRO(
+    fromModeKey: bigint,
+    toDescriptor: PipelineStateDescriptor,
+    options: { sync?: boolean } = {},
+  ): SlotTableEntry {
+    this.removeRO(fromModeKey);
+    return this.addRO(toDescriptor, options);
+  }
+
+  /** Sync slot lookup by modeKey. Returns the entry or null. */
+  lookup(modeKey: bigint): SlotTableEntry | null {
+    return this.byKey.get(modeKey) ?? null;
+  }
+
+  /**
+   * Resolves when every currently-registered slot has a linked
+   * pipeline. Subsequent `addRO` calls after this resolves with
+   * `sync: false` will produce fresh pending work; call `ready()`
+   * again to wait.
+   */
+  async ready(): Promise<void> {
+    if (this.pending.size === 0) return;
+    // Snapshot current pending set; new ones added during await
+    // resolve on their own callers' awaits if needed.
+    const snap = [...this.pending];
+    await Promise.all(snap);
+  }
+}

package/src/runtime/derivedUniforms/marker.ts

@@ -144,24 +144,37 @@ export class DerivedExpr {
   reflect(n: DerivedExpr): DerivedExpr { return new DerivedExpr(intrinsicExpr("reflect", this.ir.type, [this.ir, n.ir])); }
 }
 
-/** The `u` passed to a `derivedUniform` builder. `u.<Name>` is a mat4 trafo leaf. */
+/** The `u` passed to a `derivedUniform` builder. `u.<Name>` is a uniform leaf — a mat4
+ *  trafo by default; its real type when the wombat-shader-vite plugin supplies one (or
+ *  when you write `u.<Name>.as("vec4")`). */
 export type DerivedScope = { readonly [name: string]: DerivedExpr };
 
-function makeScope(): DerivedScope {
+/** Per-leaf type hints, keyed by uniform name (the build-time marker fills these in). */
+export type DerivedLeafTypes = Readonly<Record<string, LeafTypeName>>;
+
+function makeScope(leafTypes?: DerivedLeafTypes): DerivedScope {
   return new Proxy({} as DerivedScope, {
     get(_t, key): DerivedExpr {
       if (typeof key !== "string") throw new Error("derivedUniform: leaf names must be strings");
-      return new DerivedExpr(uniformRef(key, mat(4)));
+      const hinted = leafTypes?.[key];
+      return new DerivedExpr(uniformRef(key, hinted ? LEAF_TYPES[hinted] : mat(4)));
     },
   });
 }
 
 /**
  * Define a derived uniform. `derivedUniform(u => u.ViewTrafo.mul(u.ModelTrafo))` makes a
- * rule that slots in wherever a uniform value goes — see this module's header.
+ * rule that slots in wherever a uniform value goes — see this module's header. The optional
+ * `leafTypes` map gives `u.<Name>` leaves their real WGSL types (the wombat-shader-vite
+ * `derivedUniform(...)` marker fills it in from your `UniformScope` declarations; without
+ * the plugin, `u.<Name>` defaults to mat4 and you use `.as("vec4")` etc. when you need
+ * another type).
  */
-export function derivedUniform<T = unknown>(build: (u: DerivedScope) => DerivedExpr): DerivedRule<T> {
-  const result = build(makeScope());
+export function derivedUniform<T = unknown>(
+  build: (u: DerivedScope) => DerivedExpr,
+  leafTypes?: DerivedLeafTypes,
+): DerivedRule<T> {
+  const result = build(makeScope(leafTypes));
   if (!(result instanceof DerivedExpr)) {
     throw new Error("derivedUniform: the builder must return a derived expression, e.g. u.ViewTrafo.mul(u.ModelTrafo)");
   }