@bloopjs/toodle 0.0.100

package/src/shaders/QuadShader.ts

@@ -0,0 +1,288 @@
+import {
+  makeShaderDataDefinitions,
+  makeStructuredView,
+  type StructuredView,
+} from "webgpu-utils";
+import { WgslReflect } from "wgsl_reflect";
+import type { SceneNode } from "../scene/SceneNode";
+import type { AssetManager } from "../textures/AssetManager";
+import { assert } from "../utils/assert";
+import {
+  attachVertexInstanceBuffer,
+  createGpuPipeline,
+  getGpuPipelineDescriptor,
+  setVertexInstanceBufferLayout,
+} from "../utils/boilerplate";
+import type { EngineUniform } from "./EngineUniform";
+import type { IShader } from "./IShader";
+import {
+  codeWithLineNumbers,
+  combineShaderCode,
+  struct2BufferLayout,
+} from "./parser";
+import { pixelArtSampler, smoothSampler } from "./samplers";
+import quadWgsl from "./wgsl/quad.wgsl";
+
+export type QuadShaderOpts = {
+  assetManager?: AssetManager;
+  blendMode?: GPUBlendState;
+};
+
+export class QuadShader implements IShader {
+  label: string;
+  code: string;
+
+  #uniformValues: StructuredView;
+  #instanceData: InstanceData;
+  #instanceIndex = 0;
+  #instanceCount;
+
+  #device: GPUDevice;
+  #pipeline: GPURenderPipeline;
+  #bindGroups: GPUBindGroup[] = [];
+  #uniformBuffer: GPUBuffer;
+
+  startFrame(device: GPUDevice, uniform: EngineUniform) {
+    this.#instanceIndex = 0;
+
+    this.#uniformValues.set(uniform);
+    this.#uniformValues.set({
+      viewProjection: uniform.viewProjectionMatrix,
+      resolution: [uniform.resolution.width, uniform.resolution.height],
+    });
+
+    // Write view projection matrix to uniform buffer
+    device.queue.writeBuffer(
+      this.#uniformBuffer,
+      0,
+      this.#uniformValues.arrayBuffer,
+    );
+  }
+
+  processBatch(renderPass: GPURenderPassEncoder, nodes: SceneNode[]) {
+    renderPass.setPipeline(this.#pipeline);
+    const batchStartInstanceIndex = this.#instanceIndex;
+
+    // Count for the number of instances in the buffer...
+    let instanceCount = 0;
+
+    if (nodes.length > this.#instanceCount) {
+      throw new Error(
+        `ToodleInstanceCap: ${nodes.length} instances enqueued, max is ${this.#instanceCount} for ${this.label} shader`,
+      );
+    }
+
+    for (let i = 0; i < nodes.length; i++) {
+      if (!this.#instanceData) {
+        continue;
+      }
+      const instance = nodes[i];
+      assert(instance.renderComponent, "instance has no render component");
+      const floatOffset =
+        ((batchStartInstanceIndex + instanceCount) *
+          this.#instanceData.bufferLayout.arrayStride) /
+        Float32Array.BYTES_PER_ELEMENT;
+
+      instanceCount += instance.renderComponent.writeInstance(
+        instance,
+        this.#instanceData.cpuBuffer,
+        floatOffset,
+      );
+    }
+
+    if (this.#instanceData) {
+      const byteOffset =
+        batchStartInstanceIndex * this.#instanceData.bufferLayout.arrayStride;
+      const byteLength =
+        instanceCount * this.#instanceData.bufferLayout.arrayStride;
+
+      this.#device.queue.writeBuffer(
+        this.#instanceData.gpuBuffer,
+        byteOffset,
+        this.#instanceData.cpuBuffer,
+        byteOffset / Float32Array.BYTES_PER_ELEMENT,
+        byteLength / Float32Array.BYTES_PER_ELEMENT,
+      );
+
+      attachVertexInstanceBuffer(renderPass, this.#instanceData.gpuBuffer);
+
+      for (let i = 0; i < this.#bindGroups.length; i++) {
+        renderPass.setBindGroup(i, this.#bindGroups[i]);
+      }
+    }
+
+    this.#instanceIndex += instanceCount;
+
+    renderPass.draw(4, instanceCount, 0, batchStartInstanceIndex);
+    return 1;
+  }
+
+  endFrame() {}
+
+  constructor(
+    label: string,
+    assetManager: AssetManager,
+    device: GPUDevice,
+    presentationFormat: GPUTextureFormat,
+    userCode: string,
+    instanceCount: number,
+    blendMode?: GPUBlendState,
+  ) {
+    this.label = label;
+
+    // Combine user code with base quad shader code
+    const shaderDescriptor = combineShaderCode(label, quadWgsl, userCode);
+
+    // Create shader module from combined code
+    let module: GPUShaderModule;
+    try {
+      module = device.createShaderModule({
+        label,
+        code: shaderDescriptor.code,
+      });
+    } catch (e) {
+      console.error(codeWithLineNumbers(shaderDescriptor.code));
+      throw e;
+    }
+
+    // Store combined code for debugging
+    this.code = shaderDescriptor.code;
+
+    // Create blend state
+    const blend = blendMode ?? {
+      color: {
+        srcFactor: "src-alpha",
+        dstFactor: "one-minus-src-alpha",
+        operation: "add",
+      },
+      alpha: {
+        srcFactor: "one",
+        dstFactor: "one-minus-src-alpha",
+        operation: "add",
+      },
+    };
+
+    // Create instance data from shader code
+    const ast = new WgslReflect(shaderDescriptor.code);
+    const instanceStruct = ast.structs.find((s) =>
+      s.name.endsWith("InstanceData"),
+    );
+    if (!instanceStruct) {
+      console.error(codeWithLineNumbers(shaderDescriptor.code));
+      throw new Error(
+        "Quad shader has no instance struct. The wgsl is required to contain a struct ending in InstanceData.",
+      );
+    }
+    const bufferLayout = struct2BufferLayout(instanceStruct);
+    this.#instanceData = {
+      bufferLayout,
+      cpuBuffer: new Float32Array(
+        (bufferLayout.arrayStride * instanceCount) /
+          Float32Array.BYTES_PER_ELEMENT,
+      ),
+      gpuBuffer: device.createBuffer({
+        size: bufferLayout.arrayStride * instanceCount,
+        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
+        label: `${label} instance data`,
+      }),
+    };
+
+    // Create pipeline descriptor
+    const pipelineDescriptor = getGpuPipelineDescriptor(
+      shaderDescriptor,
+      module,
+      presentationFormat,
+      blend,
+    );
+
+    // Set vertex instance buffer layout
+    setVertexInstanceBufferLayout(
+      pipelineDescriptor,
+      this.#instanceData.bufferLayout,
+    );
+
+    // Create pipeline
+    this.#pipeline = createGpuPipeline(
+      device,
+      shaderDescriptor,
+      pipelineDescriptor,
+    );
+
+    // Reference texture atlas
+    const textureAtlas = assetManager.textureAtlas;
+
+    // Store device
+    this.#device = device;
+
+    // Create uniform buffer for engine uniforms (mat3x3 viewProjection)
+    const defs = makeShaderDataDefinitions(this.code);
+    this.#uniformValues = makeStructuredView(defs.uniforms.engineUniform);
+    this.#uniformBuffer = device.createBuffer({
+      label: `${label} engine uniform buffer`,
+      size: this.#uniformValues.arrayBuffer.byteLength,
+      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
+    });
+
+    // Create bind groups
+    this.#bindGroups = setQuadBindGroups(
+      label,
+      device,
+      this.#pipeline,
+      textureAtlas,
+      this.#uniformBuffer,
+    );
+
+    this.#instanceCount = instanceCount;
+  }
+}
+
+function setQuadBindGroups(
+  label: string,
+  device: GPUDevice,
+  pipeline: GPURenderPipeline,
+  textureAtlas: GPUTexture,
+  buffer: GPUBuffer,
+) {
+  const bindGroup = device.createBindGroup({
+    label: `${label} engine bind group`,
+    layout: pipeline.getBindGroupLayout(0),
+    entries: [
+      {
+        binding: 0,
+        resource: {
+          buffer,
+        },
+      },
+      {
+        binding: 1,
+        resource: device.createSampler(smoothSampler),
+      },
+      {
+        binding: 2,
+        resource: device.createSampler(pixelArtSampler),
+      },
+    ],
+  });
+
+  const atlasBindGroup = device.createBindGroup({
+    label: `${label} atlas bind group`,
+    layout: pipeline.getBindGroupLayout(1),
+    entries: [
+      {
+        binding: 0,
+        resource: textureAtlas.createView({
+          dimension: "2d-array",
+          arrayLayerCount: textureAtlas.depthOrArrayLayers,
+        }),
+      },
+    ],
+  });
+
+  return [bindGroup, atlasBindGroup];
+}
+
+type InstanceData = {
+  cpuBuffer: Float32Array<ArrayBuffer>;
+  gpuBuffer: GPUBuffer;
+  bufferLayout: GPUVertexBufferLayout;
+};

package/src/shaders/ShaderDescriptor.ts

@@ -0,0 +1,6 @@
+export type ShaderDescriptor = {
+  label: string;
+  code: string;
+  vertexEntrypoint: string;
+  fragmentEntrypoint: string;
+};

package/src/shaders/mod.ts

@@ -0,0 +1,5 @@
+export * from "./EngineUniform";
+export * from "./IShader";
+export * from "./postprocess/mod";
+export * from "./ShaderDescriptor";
+export * from "./samplers";

package/src/shaders/parser.ts

@@ -0,0 +1,221 @@
+import { type StructInfo, WgslReflect } from "wgsl_reflect";
+import type { ShaderDescriptor } from "./ShaderDescriptor";
+
+export function codeWithLineNumbers(code: string) {
+  return code
+    .split("\n")
+    .map((line, index) => `${index + 1}: ${line}`)
+    .join("\n");
+}
+
+export function combineShaderCode(
+  label: string,
+  base: string,
+  mod: string,
+): ShaderDescriptor {
+  let baseAst: WgslReflect;
+  let modAst: WgslReflect;
+
+  try {
+    baseAst = new WgslReflect(base);
+  } catch (e) {
+    console.error(codeWithLineNumbers(base));
+    throw e;
+  }
+  try {
+    modAst = new WgslReflect(mod);
+  } catch (e) {
+    console.error(codeWithLineNumbers(mod));
+    throw e;
+  }
+
+  if (modAst.structs.length > 1) {
+    throw new Error(
+      "Shader has more than one struct. Only one struct is supported for now.",
+    );
+  }
+
+  const hasInstanceStruct = modAst.structs.length;
+  const hasFragmentEntrypoint = modAst.entry.fragment.length > 0;
+  const hasVertexEntrypoint = modAst.entry.vertex.length > 0;
+
+  const fragmentEntrypoint = hasFragmentEntrypoint
+    ? modAst.entry.fragment[0]?.name
+    : baseAst.entry.fragment[0]?.name;
+  const vertexEntrypoint = hasVertexEntrypoint
+    ? modAst.entry.vertex[0]?.name
+    : baseAst.entry.vertex[0]?.name;
+
+  if (hasInstanceStruct) {
+    base = injectInstanceData(mod, base);
+  }
+
+  return {
+    label,
+    code: `${base}\n//==========\n\n${mod.trim()}`,
+    vertexEntrypoint,
+    fragmentEntrypoint,
+  };
+}
+
+function injectInstanceData(instanceStruct: string, code: string) {
+  const codeAst = new WgslReflect(code);
+  const snippetAst = new WgslReflect(instanceStruct);
+
+  const baseInstanceStruct = codeAst.structs.find((s) =>
+    s.name.endsWith("InstanceData"),
+  );
+  if (!baseInstanceStruct) {
+    throw new Error(
+      `No base instance struct found in code. Looking for a struct named ending in "InstanceData". Code:\n${code}`,
+    );
+  }
+
+  const baseFragmentStruct = codeAst.structs.find((s) =>
+    s.name.endsWith("VertexOutput"),
+  );
+  if (!baseFragmentStruct) {
+    throw new Error(
+      `No base fragment struct found in code. Looking for a struct named ending in "VertexOutput". Code:\n${code}`,
+    );
+  }
+
+  const struct = snippetAst.structs[0];
+  const slug = struct.name.toLowerCase();
+
+  const passthroughSnippet = struct.members
+    .map((m) => `output.${slug}_${m.name} = instance.${slug}_${m.name};`)
+    .join("\n");
+
+  const prefixedMembersSnippet = struct.members.map(
+    (m) =>
+      `${slug}_${m.name}: ${m.type.name}${m.format ? `<${m.format.name}>` : ""}`,
+  );
+
+  let startingVertexLocation = findStartingLocation(baseInstanceStruct);
+
+  const instanceInputSnippet = prefixedMembersSnippet
+    .map((snippet) => `@location(${startingVertexLocation++}) ${snippet},`)
+    .join("\n  ");
+
+  let startingFragmentLocation = findStartingLocation(baseFragmentStruct);
+
+  const vertexOutputSnippet = prefixedMembersSnippet
+    .map(
+      (snippet) =>
+        `@location(${startingFragmentLocation++}) @interpolate(flat) ${snippet},`,
+    )
+    .join("\n  ");
+
+  return code
+    .replace(/\/\/ @INSTANCE_DATA SNIPPET/, instanceInputSnippet)
+    .replace(/\/\/ @PASSTHROUGH_SNIPPET/, passthroughSnippet)
+    .replace(/\/\/ @VERTEX_OUTPUT SNIPPET/, vertexOutputSnippet);
+}
+
+function findStartingLocation(struct: StructInfo) {
+  let startingLocation = 0;
+  for (const member of struct.members) {
+    const locationAttr = (member.attributes || []).find(
+      (attr) => attr.name === "location",
+    );
+    if (locationAttr) {
+      const location = Number.parseInt(locationAttr.value as string);
+      startingLocation = Math.max(startingLocation, location);
+    }
+  }
+  return startingLocation + 1;
+}
+
+export function struct2BufferLayout(
+  struct: StructInfo,
+  stepMode: GPUVertexStepMode = "instance",
+): GPUVertexBufferLayout {
+  const hasBuiltin = struct.members.some((m) =>
+    m.attributes?.some((a) => a.name === "builtin"),
+  );
+  if (hasBuiltin) {
+    throw new Error(
+      "Can't generate buffer layout from struct with builtin attributes - they count towards size",
+    );
+  }
+
+  const unprocessedMembers = struct.members.filter(
+    (m) => !m.attributes || !m.attributes.some((a) => a.name === "location"),
+  );
+  for (const member of unprocessedMembers) {
+    console.error(
+      `Unprocessed member in struct ${struct.name}: ${member.name}`,
+    );
+  }
+
+  const membersWithAttributes = struct.members.filter((m) =>
+    m.attributes?.some((a) => a.name === "location"),
+  );
+
+  return {
+    arrayStride: struct.size,
+    attributes: membersWithAttributes.map((m) => {
+      const location = m.attributes!.find((a) => a.name === "location")?.value;
+      if (!location) {
+        throw new Error(`Location attribute not found for member: ${m.name}`);
+      }
+      if (Array.isArray(location)) {
+        throw new Error(`Location attribute is an array for member: ${m.name}`);
+      }
+
+      const shaderLocation = Number.parseInt(location);
+      if (Number.isNaN(shaderLocation)) {
+        throw new Error(
+          `Invalid location attribute: ${location} for member: ${m.name}`,
+        );
+      }
+
+      return {
+        shaderLocation,
+        offset: m.offset,
+
+        format: getGpuFormat(m.type.name),
+      } satisfies GPUVertexAttribute;
+    }),
+    stepMode,
+  } satisfies GPUVertexBufferLayout;
+}
+
+export function validateFragmentShader(shaderCode: string) {
+  const ast = new WgslReflect(shaderCode);
+  const exampleSnippet = `
+@fragment
+fn myFragmentShader(vertex: VertexOutput) -> @location(0) vec4<f32> {
+  let DEFAULT_COLOR = default_fragment_shader(vertex, nearestSampler);
+  return DEFAULT_COLOR;
+}
+`;
+
+  if (!ast.entry.fragment.length) {
+    throw new Error(`Shader code must have @fragment, e.g. ${exampleSnippet}`);
+  }
+}
+
+export function getGpuFormat(typeName: string): GPUVertexFormat {
+  // Remove <f32> or <u32> from type names
+  const baseType = typeName.replace(/<[fu]32>/, "");
+
+  const formatMap: Record<string, GPUVertexFormat> = {
+    vec2: "float32x2",
+    vec2f: "float32x2",
+    vec3: "float32x3",
+    vec3f: "float32x3",
+    vec4: "float32x4",
+    vec4f: "float32x4",
+    mat3x3: "float32x3",
+    u32: "uint32",
+    f32: "float32",
+  };
+
+  const format = formatMap[baseType.toLowerCase()];
+  if (!format) {
+    throw new Error(`Unsupported type: ${typeName}`);
+  }
+  return format;
+}