@aethervtt/shard-shader 0.0.1 → 0.1.0

package/dist/baseline/wgsl.js

@@ -0,0 +1,355 @@
+function tokenize(code) {
+  const out = [];
+  let i = 0;
+  const n = code.length;
+  while (i < n) {
+    const c = code.charCodeAt(i);
+    if (c === 32 || c === 9 || c === 10 || c === 13) {
+      i++;
+      continue;
+    }
+    if (c === 47 && code.charCodeAt(i + 1) === 47) {
+      while (i < n && code.charCodeAt(i) !== 10) i++;
+      continue;
+    }
+    if (c === 47 && code.charCodeAt(i + 1) === 42) {
+      let depth = 1;
+      i += 2;
+      while (i < n && depth > 0) {
+        if (code.charCodeAt(i) === 47 && code.charCodeAt(i + 1) === 42) {
+          depth++;
+          i += 2;
+        } else if (code.charCodeAt(i) === 42 && code.charCodeAt(i + 1) === 47) {
+          depth--;
+          i += 2;
+        } else i++;
+      }
+      continue;
+    }
+    const start = i;
+    if (isIdentStart(c)) {
+      while (i < n && isIdentPart(code.charCodeAt(i))) i++;
+      out.push({ kind: "ident", text: code.slice(start, i), start, end: i });
+      continue;
+    }
+    if (c >= 48 && c <= 57) {
+      while (i < n && /[0-9a-fA-FxXuif.eEpP+-]/.test(code[i])) {
+        const ch = code[i];
+        if ((ch === "+" || ch === "-") && !/[eEpP]/.test(code[i - 1])) break;
+        i++;
+      }
+      out.push({ kind: "number", text: code.slice(start, i), start, end: i });
+      continue;
+    }
+    if (c === 45 && code.charCodeAt(i + 1) === 62) {
+      out.push({ kind: "punct", text: "->", start, end: i + 2 });
+      i += 2;
+      continue;
+    }
+    out.push({ kind: "punct", text: code[i], start, end: i + 1 });
+    i++;
+  }
+  return out;
+}
+function isIdentStart(c) {
+  return c >= 65 && c <= 90 || c >= 97 && c <= 122 || c === 95;
+}
+function isIdentPart(c) {
+  return isIdentStart(c) || c >= 48 && c <= 57;
+}
+function matching(tokens, open) {
+  const o = tokens[open].text;
+  const c = o === "(" ? ")" : o === "[" ? "]" : o === "{" ? "}" : ">";
+  let depth = 0;
+  for (let i = open; i < tokens.length; i++) {
+    const t = tokens[i].text;
+    if (t === o) depth++;
+    else if (t === c) {
+      depth--;
+      if (depth === 0) return i;
+    }
+  }
+  return -1;
+}
+function declarations(code, tokens) {
+  const out = [];
+  let i = 0;
+  while (i < tokens.length) {
+    const first = i;
+    const attributes = [];
+    while (tokens[i]?.text === "@") {
+      const name2 = tokens[i + 1].text;
+      i += 2;
+      let args = "";
+      if (tokens[i]?.text === "(") {
+        const close = matching(tokens, i);
+        args = code.slice(tokens[i].end, tokens[close].start).trim();
+        i = close + 1;
+      }
+      attributes.push({ name: name2, args });
+    }
+    const head = tokens[i];
+    if (!head) break;
+    const kw = head.text;
+    const keyword = kw === "struct" || kw === "fn" || kw === "var" || kw === "const" || kw === "override" || kw === "alias" ? kw : "other";
+    i++;
+    let template = "";
+    if (keyword === "var" && tokens[i]?.text === "<") {
+      const close = matching(tokens, i);
+      template = code.slice(tokens[i].end, tokens[close].start).replace(/\s+/g, " ").trim();
+      i = close + 1;
+    }
+    const name = tokens[i]?.kind === "ident" ? tokens[i].text : "";
+    let type = "";
+    let last = i;
+    if (keyword === "struct" || keyword === "fn") {
+      while (last < tokens.length && tokens[last].text !== "{") last++;
+      last = matching(tokens, last);
+      if (tokens[last + 1]?.text === ";") last++;
+    } else {
+      let depth = 0;
+      let colon = -1;
+      while (last < tokens.length) {
+        const t = tokens[last].text;
+        if (t === "(" || t === "[" || t === "{") depth++;
+        else if (t === ")" || t === "]" || t === "}") depth--;
+        else if (t === ":" && depth === 0 && colon === -1) colon = last;
+        else if (t === ";" && depth === 0) break;
+        else if (t === "=" && depth === 0 && colon !== -1 && !type) {
+          type = code.slice(tokens[colon].end, tokens[last].start).trim();
+        }
+        last++;
+      }
+      if (colon !== -1 && !type) type = code.slice(tokens[colon].end, tokens[last].start).trim();
+    }
+    const end = tokens[Math.min(last, tokens.length - 1)].end;
+    out.push({
+      attributes,
+      keyword,
+      name,
+      template,
+      type,
+      start: tokens[first].start,
+      end,
+      first,
+      last: Math.min(last, tokens.length - 1)
+    });
+    i = last + 1;
+  }
+  return out;
+}
+function attribute(d, name) {
+  return d.attributes.find((a) => a.name === name)?.args;
+}
+const SHORT = { f: "f32", u: "u32", i: "i32", h: "f16" };
+function parseType(text) {
+  const t = text.replace(/\s+/g, "");
+  if (t === "f32" || t === "u32" || t === "i32" || t === "f16" || t === "bool") {
+    return { kind: "scalar", name: t };
+  }
+  let m = /^vec([234])(?:([fuih])|<(\w+)>)$/.exec(t);
+  if (m) {
+    const element = parseType(m[2] ? SHORT[m[2]] : m[3]);
+    return { kind: "vector", size: Number(m[1]), element };
+  }
+  m = /^mat([234])x([234])(?:([fh])|<(\w+)>)$/.exec(t);
+  if (m) {
+    const element = parseType(m[3] ? SHORT[m[3]] : m[4]);
+    return {
+      kind: "matrix",
+      columns: Number(m[1]),
+      rows: Number(m[2]),
+      element
+    };
+  }
+  if (t.startsWith("array<") && t.endsWith(">")) {
+    const inner = t.slice(6, -1);
+    let depth = 0;
+    let split = -1;
+    for (let i = 0; i < inner.length; i++) {
+      const c = inner[i];
+      if (c === "<") depth++;
+      else if (c === ">") depth--;
+      else if (c === "," && depth === 0) split = i;
+    }
+    const element = parseType(split === -1 ? inner : inner.slice(0, split));
+    const countText = split === -1 ? void 0 : inner.slice(split + 1).replace(/[ui]$/, "");
+    const count = countText === void 0 ? void 0 : Number(countText);
+    return {
+      kind: "array",
+      element,
+      elementText: split === -1 ? inner : inner.slice(0, split),
+      count: Number.isFinite(count) ? count : void 0,
+      text
+    };
+  }
+  if (t.startsWith("atomic<") && t.endsWith(">")) {
+    return { kind: "atomic", element: parseType(t.slice(7, -1)) };
+  }
+  if (/^\w+$/.test(t)) return { kind: "struct", name: t };
+  return { kind: "other", text };
+}
+const roundUp = (n, a) => Math.ceil(n / a) * a;
+class Layouts {
+  structs = /* @__PURE__ */ new Map();
+  cache = /* @__PURE__ */ new Map();
+  constructor(code, tokens, decls) {
+    for (const d of decls) {
+      if (d.keyword !== "struct") continue;
+      let i = d.first;
+      while (tokens[i].text !== "{") i++;
+      const close = matching(tokens, i);
+      const members = [];
+      i++;
+      while (i < close) {
+        let align;
+        let size;
+        while (tokens[i].text === "@") {
+          const attr = tokens[i + 1].text;
+          i += 2;
+          if (tokens[i].text === "(") {
+            const end = matching(tokens, i);
+            const value = Number(
+              code.slice(tokens[i].end, tokens[end].start).trim().replace(/[ui]$/, "")
+            );
+            if (attr === "align") align = value;
+            if (attr === "size") size = value;
+            i = end + 1;
+          }
+        }
+        const name = tokens[i].text;
+        i += 2;
+        let depth = 0;
+        const from = i;
+        while (i < close) {
+          const t = tokens[i].text;
+          if (t === "<" || t === "(") depth++;
+          else if (t === ">" || t === ")") depth--;
+          else if (t === "," && depth === 0) break;
+          i++;
+        }
+        members.push({
+          name,
+          typeText: code.slice(tokens[from].start, tokens[i - 1].end),
+          align,
+          size
+        });
+        if (tokens[i]?.text === ",") i++;
+      }
+      this.structs.set(d.name, { name: d.name, members });
+    }
+  }
+  has(name) {
+    return this.structs.has(name);
+  }
+  struct(name, space) {
+    const key = `${space}:${name}`;
+    const cached = this.cache.get(key);
+    if (cached) return cached;
+    const s = this.structs.get(name);
+    if (!s) throw new Error(`No struct ${name}`);
+    let offset = 0;
+    let align = 1;
+    const members = [];
+    for (const m of s.members) {
+      const type = parseType(m.typeText);
+      const a = m.align ?? this.align(type, space);
+      const size = m.size ?? this.size(type, space);
+      offset = roundUp(offset, a);
+      members.push({ name: m.name, type, typeText: m.typeText, offset, size, align: a });
+      offset += size;
+      align = Math.max(align, a);
+    }
+    if (space === "uniform") align = roundUp(align, 16);
+    const layout = { name, members, size: roundUp(offset, align), align };
+    this.cache.set(key, layout);
+    return layout;
+  }
+  align(t, space) {
+    switch (t.kind) {
+      case "scalar":
+        return t.name === "f16" ? 2 : 4;
+      case "atomic":
+        return 4;
+      case "vector": {
+        const e = t.element.name === "f16" ? 2 : 4;
+        return t.size === 2 ? 2 * e : 4 * e;
+      }
+      case "matrix":
+        return this.align({ kind: "vector", size: t.rows, element: t.element }, space);
+      case "array": {
+        const a = this.align(t.element, space);
+        return space === "uniform" ? roundUp(a, 16) : a;
+      }
+      case "struct":
+        return this.struct(t.name, space).align;
+      default:
+        throw new Error(`No layout for ${t.text}`);
+    }
+  }
+  size(t, space) {
+    switch (t.kind) {
+      case "scalar":
+        return t.name === "f16" ? 2 : 4;
+      case "atomic":
+        return 4;
+      case "vector":
+        return t.size * (t.element.name === "f16" ? 2 : 4);
+      case "matrix": {
+        const column = { kind: "vector", size: t.rows, element: t.element };
+        return t.columns * roundUp(this.size(column, space), this.align(column, space));
+      }
+      case "array":
+        if (t.count === void 0) return this.stride(t, space);
+        return t.count * this.stride(t, space);
+      case "struct":
+        return this.struct(t.name, space).size;
+      default:
+        throw new Error(`No layout for ${t.text}`);
+    }
+  }
+  /** Bytes between an array's elements. */
+  stride(t, space) {
+    const e = t.element;
+    const stride = roundUp(this.size(e, space), this.align(e, space));
+    return space === "uniform" ? roundUp(stride, 16) : stride;
+  }
+  /**
+   * Whether a type lies out the same in uniform and storage memory: every member and element at
+   * the same offset, so the same bytes read the same through either.
+   */
+  sameInUniform(t) {
+    try {
+      return this.sameLayout(t);
+    } catch {
+      return false;
+    }
+  }
+  sameLayout(t) {
+    switch (t.kind) {
+      case "scalar":
+      case "vector":
+      case "matrix":
+        return true;
+      case "array":
+        return t.count !== void 0 && this.stride(t, "storage") === this.stride(t, "uniform") && this.sameLayout(t.element);
+      case "struct": {
+        const s = this.struct(t.name, "storage");
+        const u = this.struct(t.name, "uniform");
+        return s.members.every(
+          (m, i) => m.offset === u.members[i].offset && this.sameLayout(m.type)
+        );
+      }
+      default:
+        return false;
+    }
+  }
+}
+export {
+  Layouts,
+  attribute,
+  declarations,
+  matching,
+  parseType,
+  tokenize
+};

package/dist/data-marks.d.ts

@@ -0,0 +1,23 @@
+/** Bumps when the baseline rewrite's output changes: baked baseline code from before goes stale. */
+export declare const BASELINE_REWRITE_VERSION = 1;
+/**
+ * Width of a data texture in texels (0064): element bytes at byte offset b sit in texel b / 16, a
+ * row holding this many. The rewrite's loaders and the engine's data textures agree on it.
+ */
+export declare const DATA_TEXTURE_WIDTH = 1024;
+/** A storage declaration marked `@data`, as collected before linking. */
+export interface DataDeclaration {
+    /** The module that declared it, for errors. */
+    module: string;
+    name: string;
+    group: number;
+    binding: number;
+    kind: 'texture' | 'uniform';
+    /** `@data(uniform, N)`: the uniform array length of a runtime-sized array. */
+    count?: number;
+}
+/** The source without its `@data` marks, and the declarations they marked. */
+export declare function collectData(source: string, module: string): {
+    code: string;
+    data: DataDeclaration[];
+};

package/dist/data-marks.js

@@ -0,0 +1,40 @@
+import { ShardError } from "@aethervtt/shard-core";
+const BASELINE_REWRITE_VERSION = 1;
+const DATA_TEXTURE_WIDTH = 1024;
+const MARK = /@data(?:\((texture|uniform)(?:\s*,\s*(\d+))?\))?\s/g;
+const DECLARATION = /^((?:@\w+\s*(?:\([^)]*\))?\s*)*)var\s*<\s*storage[^>]*>\s*([A-Za-z_]\w*)/;
+function collectData(source, module) {
+  if (!source.includes("@data")) return { code: source, data: [] };
+  const data = [];
+  let code = "";
+  let at = 0;
+  for (const m of source.matchAll(MARK)) {
+    const after = source.slice(m.index + m[0].length);
+    const d = DECLARATION.exec(after);
+    const group = d ? /@group\s*\(\s*(\d+)/.exec(d[1]) : null;
+    const binding = d ? /@binding\s*\(\s*(\d+)/.exec(d[1]) : null;
+    if (!d || !group || !binding) {
+      throw new ShardError(
+        "shader/data-misplaced",
+        `@data in ${module} must mark a storage declaration with @group and @binding`,
+        { hint: "@data @group(2) @binding(0) var<storage, read> name: array<T>;" }
+      );
+    }
+    data.push({
+      module,
+      name: d[2],
+      group: Number(group[1]),
+      binding: Number(binding[1]),
+      kind: m[1] === "uniform" ? "uniform" : "texture",
+      ...m[2] ? { count: Number(m[2]) } : {}
+    });
+    code += source.slice(at, m.index);
+    at = m.index + m[0].length;
+  }
+  return { code: code + source.slice(at), data };
+}
+export {
+  BASELINE_REWRITE_VERSION,
+  DATA_TEXTURE_WIDTH,
+  collectData
+};

package/dist/index.d.ts

@@ -1,3 +1,4 @@
+export { DATA_TEXTURE_WIDTH } from './data-marks.js';
 export { applyHooks, findHooks, type HookSignature } from './hooks.js';
 export { type WgslLayout, wgslLayout } from './layout.js';
 export { type BakedShader, type LinkedShader, type LinkRequest, type ShaderBake, ShaderLibrary, type SourceLocation, } from './library.js';

package/dist/index.js

@@ -1,9 +1,11 @@
+import { DATA_TEXTURE_WIDTH } from "./data-marks.js";
 import { applyHooks, findHooks } from "./hooks.js";
 import { wgslLayout } from "./layout.js";
 import {
   ShaderLibrary
 } from "./library.js";
 export {
+  DATA_TEXTURE_WIDTH,
   ShaderLibrary,
   applyHooks,
   findHooks,

package/dist/layout.d.ts

@@ -6,6 +6,8 @@ interface FieldLayout {
     offset: number;
     size: number;
     align: number;
+    /** Struct fields: the nested struct's members. */
+    members?: readonly FieldLayout[];
 }
 export interface WgslLayout<F extends Fields> {
     /** WGSL struct declaration, e.g. `struct ToonParams { bands: u32, rim: vec4f, width: f32, }`. */
@@ -15,6 +17,11 @@ export interface WgslLayout<F extends Fields> {
     readonly size: number;
     readonly align: number;
     readonly fields: readonly FieldLayout[];
+    /**
+     * The nested structs it uses (struct fields with a `wgsl` option), as `module::path::Name`: a
+     * module declaring this struct imports them.
+     */
+    readonly imports: readonly string[];
     /** Writes a value at a byte offset using WGSL's layout rules. */
     write(view: DataView, offset: number, value: InferFields<F>): void;
 }

package/dist/layout.js

@@ -21,82 +21,118 @@ const TYPES = {
   affine3x4: { type: "array<vec4f, 3>", size: 48, align: 16 }
 };
 const alignTo = (n, a) => Math.ceil(n / a) * a;
-function wgslLayout(def) {
-  const structName = def.name.split("/").pop();
+function pack(owner, entries, imports) {
   const fields = [];
   let offset = 0;
   let structAlign = 4;
-  for (const { name, field, storage } of def.layout) {
-    if (storage === "object") continue;
-    const t = TYPES[field.kind];
-    if (!t) {
-      throw new ShardError(
-        "shader/unsupported-field",
-        `Field "${name}" of ${def.name} (${field.kind}) can't be a GPU struct member`,
-        {
-          hint: "GPU structs take numbers, bools, enums, vectors, colors, and matrices. f64 fields are not allowed; object fields are skipped."
-        }
-      );
+  for (const [name, field] of entries) {
+    if (field.options.gpu === false) continue;
+    let layout;
+    if (field.kind === "struct" && field.options.wgsl) {
+      const inner = pack(`${owner}.${name}`, Object.entries(field.fields ?? {}), imports);
+      const path = field.options.wgsl;
+      imports.add(path);
+      const align = Math.max(16, inner.align);
+      offset = alignTo(offset, align);
+      layout = {
+        name,
+        kind: "struct",
+        wgslType: path.slice(path.lastIndexOf("::") + 2),
+        offset,
+        size: alignTo(inner.size, align),
+        align,
+        members: inner.fields
+      };
+    } else {
+      if (field.storage === "object") continue;
+      const t = TYPES[field.kind];
+      if (!t) {
+        throw new ShardError(
+          "shader/unsupported-field",
+          `Field "${name}" of ${owner} (${field.kind}) can't be a GPU struct member`,
+          {
+            hint: "GPU structs take numbers, bools, enums, vectors, colors, matrices, and structs with a `wgsl` option. f64 fields are not allowed; other object fields are skipped."
+          }
+        );
+      }
+      offset = alignTo(offset, t.align);
+      layout = { name, kind: field.kind, wgslType: t.type, offset, size: t.size, align: t.align };
     }
-    offset = alignTo(offset, t.align);
-    fields.push({ name, kind: field.kind, wgslType: t.type, offset, size: t.size, align: t.align });
-    offset += t.size;
-    structAlign = Math.max(structAlign, t.align);
+    fields.push(layout);
+    offset += layout.size;
+    structAlign = Math.max(structAlign, layout.align);
   }
-  const size = alignTo(Math.max(offset, 4), structAlign);
+  return { fields, size: alignTo(Math.max(offset, 4), structAlign), align: structAlign, imports };
+}
+function enumsOf(entries, out) {
+  for (const [, field] of entries) {
+    if (field.kind === "enum") out.set(field, field.jsonSchema().enum);
+    else if (field.kind === "struct" && field.fields) enumsOf(Object.entries(field.fields), out);
+  }
+  return out;
+}
+function wgslLayout(def) {
+  const structName = def.name.split("/").pop();
+  const entries = def.layout.map(({ name, field }) => [name, field]);
+  const packed = pack(def.name, entries, /* @__PURE__ */ new Set());
   const wgsl = `struct ${structName} {
-${fields.map((f) => `  ${f.name}: ${f.wgslType},`).join("\n")}
+${packed.fields.map((f) => `  ${f.name}: ${f.wgslType},`).join("\n")}
 }`;
-  const enumValues = /* @__PURE__ */ new Map();
-  for (const { name, field } of def.layout) {
-    if (field.kind === "enum") {
-      enumValues.set(name, field.jsonSchema().enum);
+  const schema = Object.fromEntries(entries);
+  const enumValues = enumsOf(entries, /* @__PURE__ */ new Map());
+  const writeFields = (view, base, layouts, record, schema2) => {
+    for (const f of layouts) {
+      const o = base + f.offset;
+      const v = record[f.name];
+      switch (f.kind) {
+        case "struct": {
+          writeFields(view, o, f.members, v, schema2[f.name].fields);
+          break;
+        }
+        case "f32":
+          view.setFloat32(o, v, true);
+          break;
+        case "i8":
+        case "i16":
+        case "i32":
+          view.setInt32(o, v, true);
+          break;
+        case "u8":
+        case "u16":
+        case "u32":
+          view.setUint32(o, v, true);
+          break;
+        case "bool":
+          view.setUint32(o, v ? 1 : 0, true);
+          break;
+        case "enum": {
+          const values = enumValues.get(schema2[f.name]);
+          view.setUint32(o, Math.max(0, values.indexOf(v)), true);
+          break;
+        }
+        case "mat3": {
+          const m = v;
+          for (let c = 0; c < 3; c++) {
+            for (let r = 0; r < 3; r++) view.setFloat32(o + c * 16 + r * 4, m[c * 3 + r], true);
+          }
+          break;
+        }
+        default: {
+          const arr = v;
+          for (let i = 0; i < arr.length; i++) view.setFloat32(o + i * 4, arr[i], true);
+        }
+      }
     }
-  }
+  };
   return {
     wgsl,
     structName,
-    size,
-    align: structAlign,
-    fields,
+    size: packed.size,
+    align: packed.align,
+    fields: packed.fields,
+    imports: [...packed.imports],
     write(view, base, value) {
-      const record = value;
-      for (const f of fields) {
-        const o = base + f.offset;
-        const v = record[f.name];
-        switch (f.kind) {
-          case "f32":
-            view.setFloat32(o, v, true);
-            break;
-          case "i8":
-          case "i16":
-          case "i32":
-            view.setInt32(o, v, true);
-            break;
-          case "u8":
-          case "u16":
-          case "u32":
-            view.setUint32(o, v, true);
-            break;
-          case "bool":
-            view.setUint32(o, v ? 1 : 0, true);
-            break;
-          case "enum":
-            view.setUint32(o, Math.max(0, enumValues.get(f.name).indexOf(v)), true);
-            break;
-          case "mat3": {
-            const m = v;
-            for (let c = 0; c < 3; c++) {
-              for (let r = 0; r < 3; r++) view.setFloat32(o + c * 16 + r * 4, m[c * 3 + r], true);
-            }
-            break;
-          }
-          default: {
-            const arr = v;
-            for (let i = 0; i < arr.length; i++) view.setFloat32(o + i * 4, arr[i], true);
-          }
-        }
-      }
+      writeFields(view, base, packed.fields, value, schema);
     }
   };
 }