《游戏引擎实现原理》11.2 反射驱动的参数绑定

// file: src/render/shader/uniform-reflect.ts
// tier: L1
// verify: npx vitest run src/render/shader/uniform-reflect.test.ts
//
// helio 渲染卷：从 WGSL 的 struct 反射出 uniform 内存布局，并把 JS 值按该布局打包。
// 让「参数怎么放」由着色器自己决定：C++/TS 侧不再手填偏移，改着色器里的 struct
// 就会自动反映到打包结果上（这是 11.2 的核心）。布局规则取 WGSL uniform 地址空间。

/** 支持反射的标量/向量/矩阵/数组类型。 */
export type ReflectKind = 'f32' | 'i32' | 'u32';

export interface MemberLayout {
  readonly name: string;
  readonly type: string;
  /** 在缓冲区里的字节偏移。 */
  readonly offset: number;
  /** 该成员占用的字节数（含数组元素填充，不含尾部结构填充）。 */
  readonly size: number;
  readonly kind: ReflectKind;
  /** 分量个数；数组为元素分量总数。 */
  readonly components: number;
  /** 是否为数组类型。 */
  readonly isArray: boolean;
}

export interface UniformLayout {
  readonly structName: string;
  readonly binding: number;
  readonly members: readonly MemberLayout[];
  /** 整个 uniform 缓冲区的字节大小（已按 16 对齐）。 */
  readonly size: number;
  /** 结构体对齐，uniform 地址空间里恒为 16 的正倍数。 */
  readonly align: number;
}

const roundUp = (align: number, v: number): number => Math.ceil(v / align) * align;

interface TypeInfo { align: number; size: number; kind: ReflectKind; components: number; isArray: boolean }

/** 解析单个类型文本（支持标量、vecN<f32>、matNxN<f32>、array<T, N>）。 */
export function typeInfo(type: string): TypeInfo {
  const t = type.replace(/\s+/g, '');
  const arr = /^array<(.+),(\d+)>$/.exec(t);
  if (arr) {
    const elem = typeInfo(arr[1]!);
    const n = Number(arr[2]);
    const stride = roundUp(16, elem.size); // uniform 里数组元素步长按 16 对齐
    return { align: 16, size: stride * n, kind: elem.kind, components: elem.components * n, isArray: true };
  }
  const m = /^mat(\d)x(\d)<f32>$/.exec(t);
  if (m) {
    const cols = Number(m[1]);
    // 每列是 vecR，按 16 对齐（vec3 列会被填到 16）
    const colSize = roundUp(16, Number(m[2]) * 4);
    return { align: 16, size: colSize * cols, kind: 'f32', components: cols * Number(m[2]), isArray: false };
  }
  const v = /^vec([234])<f32>$/.exec(t);
  if (v) {
    const n = Number(v[1]);
    return { align: n === 2 ? 8 : 16, size: n * 4, kind: 'f32', components: n, isArray: false };
  }
  if (t === 'f32' || t === 'i32' || t === 'u32') {
    return { align: 4, size: 4, kind: t, components: 1, isArray: false };
  }
  throw new Error(`不支持的反射类型: ${type}`);
}

/** 按 `<>` 深度切分 struct 成员（数组里的逗号不算分隔）。 */
function splitMembers(body: string): string[] {
  const parts: string[] = [];
  let depth = 0, cur = '';
  for (const ch of body) {
    if (ch === '<') depth++;
    else if (ch === '>') depth--;
    if (ch === ',' && depth === 0) { if (cur.trim()) parts.push(cur); cur = ''; continue; }
    cur += ch;
  }
  if (cur.trim()) parts.push(cur);
  return parts;
}

/** 反射 WGSL 源码里第一个 `var<uniform>` 所指结构体的布局。 */
export function reflectUniform(wgsl: string): UniformLayout {
  const bind = /@group\((\d+)\)\s*@binding\((\d+)\)\s*var<uniform>\s+([A-Za-z_]\w*)\s*:\s*([A-Za-z_]\w*)/.exec(wgsl);
  if (!bind) throw new Error('未找到 var<uniform> 绑定');
  const structName = bind[4]!;
  const sm = new RegExp(`struct\\s+${structName}\\s*\\{([\\s\\S]*?)\\}`).exec(wgsl);
  if (!sm) throw new Error(`未找到 struct ${structName}`);
  const members: MemberLayout[] = [];
  let offset = 0, maxAlign = 1;
  for (const raw of splitMembers(sm[1]!)) {
    const mm = /^\s*([A-Za-z_]\w*)\s*:\s*(.+?)\s*$/.exec(raw);
    if (!mm) throw new Error(`无法解析成员: ${raw}`);
    const info = typeInfo(mm[2]!);
    offset = roundUp(info.align, offset);
    members.push({
      name: mm[1]!, type: mm[2]!, offset, size: info.size,
      kind: info.kind, components: info.components, isArray: info.isArray,
    });
    offset += info.size;
    maxAlign = Math.max(maxAlign, info.align);
  }
  const align = roundUp(16, maxAlign);
  return { structName, binding: Number(bind[2]), members, size: roundUp(align, offset), align };
}

/** 按布局把值打包进 ArrayBuffer；缺字段按 0 处理。 */
export function packUniform(layout: UniformLayout, values: Readonly<Record<string, number | number[]>>): ArrayBuffer {
  const buf = new ArrayBuffer(layout.size);
  const dv = new DataView(buf);
  for (const m of layout.members) {
    const v = values[m.name];
    if (v === undefined) continue;
    const arr = Array.isArray(v) ? v : [v];
    for (let i = 0; i < m.components; i++) {
      const x = arr[i] ?? 0;
      const at = m.offset + i * 4;
      if (m.kind === 'f32') dv.setFloat32(at, x, true);
      else if (m.kind === 'i32') dv.setInt32(at, x, true);
      else dv.setUint32(at, x, true);
    }
  }
  return buf;
}

// file: shaders/material-uniform.wgsl
// tier: L2
// verify: node tools/run-l2.mjs src/render/shader/uniform-reflect-l2.ts

struct Material {
  baseColor: vec4<f32>,
  roughness: f32,
  metallic: f32,
  uvScale: vec2<f32>,
};
@group(0) @binding(0) var<uniform> mat: Material;

struct VSOut { @builtin(position) pos: vec4<f32>, @location(0) uv: vec2<f32> };

@vertex
fn vs(@builtin(vertex_index) i: u32) -> VSOut {
  var p = array<vec2<f32>, 3>(
    vec2<f32>(-1.0, -1.0), vec2<f32>(3.0, -1.0), vec2<f32>(-1.0, 3.0),
  );
  var o: VSOut;
  o.pos = vec4<f32>(p[i], 0.0, 1.0);
  o.uv = p[i] * 0.5 + vec2<f32>(0.5);
  return o;
}

@fragment
fn fs(in: VSOut) -> @location(0) vec4<f32> {
  let shade = mat.baseColor.rgb * (1.0 - mat.roughness * mat.metallic);
  let uv = in.uv * mat.uvScale;
  return vec4<f32>(shade * step(0.0, uv.x), mat.baseColor.a);
}

offsets    {baseColor: 0, roughness: 16, metallic: 20, uvScale: 24}
size       32
errors     0
center     [153, 76, 38, 255]