wgsl-test 0.2.1

package/package.json

@@ -0,0 +1,61 @@
+{
+  "name": "wgsl-test",
+  "version": "0.2.1",
+  "type": "module",
+  "files": [
+    "dist",
+    "images",
+    "src"
+  ],
+  "repository": "github:wgsl-tooling-wg/wesl-js",
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
+  "dependencies": {
+    "pngjs": "^7.0.0",
+    "thimbleberry": "^0.2.10",
+    "webgpu": "^0.3.8",
+    "wesl": "0.6.47",
+    "wesl-gpu": "0.1.1"
+  },
+  "devDependencies": {
+    "@types/pngjs": "^6.0.0",
+    "@webgpu/types": "^0.1.65",
+    "dependent_package": "x",
+    "wesl-tooling": "x"
+  },
+  "peerDependencies": {
+    "vitest": "^3.2.4",
+    "vitest-image-snapshot": "^0.6.38"
+  },
+  "peerDependenciesMeta": {
+    "vitest": {
+      "optional": true
+    },
+    "vitest-image-snapshot": {
+      "optional": true
+    }
+  },
+  "keywords": [
+    "shader",
+    "snapshot",
+    "test",
+    "testing",
+    "unit-test",
+    "visual-regression",
+    "WebGPU",
+    "WESL",
+    "WGSL"
+  ],
+  "scripts": {
+    "build": "tsdown",
+    "dev": "tsdown --watch",
+    "test": "cross-env FORCE_COLOR=1 vitest",
+    "test:once": "vitest run",
+    "typecheck": "tsgo"
+  },
+  "main": "./dist/index.js"
+}

package/src/CompileShader.ts

@@ -0,0 +1,185 @@
+import * as path from "node:path";
+import { fileURLToPath } from "node:url";
+import type { LinkParams, ModuleResolver, WeslBundle } from "wesl";
+import { CompositeResolver, link, RecordResolver } from "wesl";
+import {
+  dependencyBundles,
+  FileModuleResolver,
+  findWeslToml,
+  readPackageJson,
+  resolveProjectDir,
+} from "wesl-tooling";
+
+export interface ShaderContext {
+  /** Dependency bundles for the shader. */
+  libs: WeslBundle[];
+
+  /** Resolver for lazy loading (when useSourceShaders is true). */
+  resolver?: ModuleResolver;
+
+  /** Package name for module resolution. */
+  packageName?: string;
+}
+
+export interface ResolveContextParams {
+  /** WESL/WGSL shader source code. */
+  src: string;
+
+  /** Project directory for resolving dependencies. */
+  projectDir?: string;
+
+  /** Use source shaders instead of built bundles. Default: true. */
+  useSourceShaders?: boolean;
+}
+
+export interface CompileShaderParams {
+  /** Project directory for resolving shader dependencies.
+   * Used to locate installed npm shader libraries.
+   * Optional: defaults to searching upward from cwd for package.json or wesl.toml. */
+  projectDir?: string;
+
+  /** GPU device to use for shader compilation. */
+  device: GPUDevice;
+
+  /** WESL/WGSL shader source code to compile. */
+  src: string;
+
+  /** Conditions for conditional compilation.
+   * Used to control `@if` directives in the shader. */
+  conditions?: LinkParams["conditions"];
+
+  /** Constants for shader compilation.
+   * Injects host-provided values via the `constants::` namespace. */
+  constants?: LinkParams["constants"];
+
+  /** Virtual libraries to include in the shader.
+   * Allows dynamic generation of shader code at runtime. */
+  virtualLibs?: LinkParams["virtualLibs"];
+
+  /** Use source shaders from current package instead of built bundles.
+   * Default: true for faster iteration during development.
+   * Set to false or use TEST_BUNDLES=true environment variable to test built bundles.
+   *
+   * Precedence: explicit parameter > TEST_BUNDLES env var > default (true)
+   */
+  useSourceShaders?: boolean;
+}
+
+/**
+ * Compiles a WESL shader source string into a GPUShaderModule with automatic dependency resolution.
+ *
+ * Parses the shader source to detect references to shader packages, then automatically
+ * includes the required npm package bundles. By default, loads source shaders from the
+ * current package for fast iteration without requiring rebuilds.
+ *
+ * @returns Compiled GPUShaderModule ready for use in render or compute pipelines
+ * @throws Error if shader compilation fails with compilation error details
+ */
+export async function compileShader(
+  params: CompileShaderParams,
+): Promise<GPUShaderModule> {
+  const { device, src, conditions, constants, virtualLibs } = params;
+  const ctx = await resolveShaderContext({
+    src,
+    projectDir: params.projectDir,
+    useSourceShaders: params.useSourceShaders,
+  });
+
+  let linkParams: Pick<LinkParams, "resolver" | "libs" | "weslSrc">;
+  if (ctx.resolver) {
+    linkParams = { resolver: ctx.resolver, libs: ctx.libs };
+  } else {
+    linkParams = { weslSrc: { main: src }, libs: ctx.libs };
+  }
+
+  const linked = await link({
+    ...linkParams,
+    rootModuleName: "main",
+    virtualLibs,
+    conditions,
+    constants,
+    packageName: ctx.packageName,
+  });
+  const module = linked.createShaderModule(device);
+
+  await verifyCompilation(module);
+  return module;
+}
+
+/** Resolve project context for shader compilation: bundles, resolver, and package name. */
+export async function resolveShaderContext(
+  params: ResolveContextParams,
+): Promise<ShaderContext> {
+  const { src, useSourceShaders = !process.env.TEST_BUNDLES } = params;
+  const projectDir = await resolveProjectDir(params.projectDir);
+  const packageName = await getPackageName(projectDir);
+
+  const libs = await dependencyBundles(
+    { main: src },
+    projectDir,
+    packageName,
+    !useSourceShaders, // include current package when testing bundles
+  );
+
+  const resolver = useSourceShaders
+    ? await lazyFileResolver(projectDir, src, packageName)
+    : undefined;
+
+  return { libs, resolver, packageName };
+}
+
+/** Create a project resolver for loading modules from the filesystem.
+ * Handles wesl.toml configuration and creates FileModuleResolver with correct baseDir.
+ *
+ * @param projectDir Project directory (defaults to cwd)
+ * @param packageName Package name for module resolution (optional)
+ * @returns FileModuleResolver configured for the project
+ */
+export async function createProjectResolver(
+  projectDir?: string,
+  packageName?: string,
+): Promise<ModuleResolver> {
+  const resolved = await resolveProjectDir(projectDir);
+  const projectPath = fileURLToPath(resolved);
+  const tomlInfo = await findWeslToml(projectPath);
+  const baseDir = path.isAbsolute(tomlInfo.resolvedRoot)
+    ? tomlInfo.resolvedRoot
+    : path.join(projectPath, tomlInfo.resolvedRoot);
+
+  return new FileModuleResolver(baseDir, packageName);
+}
+
+/** Create a lazy resolver that loads local shaders on-demand from the filesystem.
+ * Laziness allows testing without rebuilding the current package after edits. */
+async function lazyFileResolver(
+  projectDir: string,
+  mainSrc: string,
+  packageName: string | undefined,
+): Promise<CompositeResolver> {
+  const mainResolver = new RecordResolver({ main: mainSrc }, { packageName });
+  const fileResolver = await createProjectResolver(projectDir, packageName);
+  return new CompositeResolver([mainResolver, fileResolver]);
+}
+
+/** Read package name from package.json, normalized for WGSL identifiers. */
+async function getPackageName(projectDir: string): Promise<string | undefined> {
+  try {
+    const pkg = await readPackageJson(projectDir);
+    const name = pkg.name as string;
+    return name.replace(/-/g, "_");
+  } catch {
+    return undefined;
+  }
+}
+
+/** Verify shader compilation succeeded, throw on errors. */
+async function verifyCompilation(module: GPUShaderModule): Promise<void> {
+  const info = await module.getCompilationInfo();
+  const errors = info.messages.filter(msg => msg.type === "error");
+  if (errors.length > 0) {
+    const messages = errors
+      .map(e => `${e.lineNum}:${e.linePos} ${e.message}`)
+      .join("\n");
+    throw new Error(`Shader compilation failed:\n${messages}`);
+  }
+}

package/src/ExampleImages.ts

@@ -0,0 +1,27 @@
+import * as path from "node:path";
+import { fileURLToPath } from "node:url";
+import { DeviceCache } from "wesl-gpu";
+import { pngToTexture } from "./ImageHelpers.ts";
+
+const textureCache = new DeviceCache<GPUTexture>();
+
+/** return a texture to the bundled lemur test image. */
+export async function lemurTexture(
+  device: GPUDevice,
+  size: 256 | 512 = 512,
+): Promise<GPUTexture> {
+  const lemurPath = lemurImagePath(size);
+  const cached = textureCache.get(device, lemurPath);
+  if (cached) return cached;
+
+  const texture = await pngToTexture(device, lemurImagePath(size));
+  textureCache.set(device, lemurPath, texture);
+
+  return texture;
+}
+
+/** Get the path to the bundled lemur test image. */
+export function lemurImagePath(size: 256 | 512 = 512): string {
+  const moduleDir = path.dirname(fileURLToPath(import.meta.url));
+  return path.join(moduleDir, "..", "images", `lemur${size}.png`);
+}

package/src/ImageHelpers.ts

@@ -0,0 +1,38 @@
+import * as fs from "node:fs/promises";
+import * as path from "node:path";
+import { PNG } from "pngjs";
+
+/** Load PNG file and create GPU texture. */
+export async function pngToTexture(
+  device: GPUDevice,
+  imagePath: string,
+): Promise<GPUTexture> {
+  const png = await loadPNG(imagePath);
+
+  const texture = device.createTexture({
+    label: `test-texture-photo-${path.basename(imagePath)}`,
+    size: { width: png.width, height: png.height, depthOrArrayLayers: 1 },
+    format: "rgba8unorm",
+    usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
+  });
+
+  device.queue.writeTexture(
+    { texture },
+    new Uint8Array(png.data),
+    { bytesPerRow: png.width * 4 },
+    { width: png.width, height: png.height },
+  );
+
+  return texture;
+}
+
+async function loadPNG(imagePath: string): Promise<PNG> {
+  const fileData = await fs.readFile(imagePath);
+  return new Promise<PNG>((resolve, reject) => {
+    const png = new PNG();
+    png.parse(fileData, (err, data) => {
+      if (err) reject(err);
+      else resolve(data);
+    });
+  });
+}

package/src/ShaderModuleLoader.ts

@@ -0,0 +1,46 @@
+import { normalizeModuleName } from "wesl";
+import { createProjectResolver } from "./CompileShader.ts";
+
+/** Validates that exactly one of src or moduleName is provided.
+ * @throws Error if neither or both are provided */
+export function validateSourceParams(src?: string, moduleName?: string): void {
+  if (!src && !moduleName) {
+    throw new Error("Either src or moduleName must be provided");
+  }
+  if (src && moduleName) {
+    throw new Error("Cannot provide both src and moduleName");
+  }
+}
+
+/** Loads shader source from module name using filesystem resolver.
+ * @param moduleName Shader module name (bare name, path, or module path)
+ * @param projectDir Project directory for module resolution
+ * @returns Shader source code
+ */
+export async function loadShaderSourceFromModule(
+  moduleName: string,
+  projectDir?: string,
+): Promise<string> {
+  const resolver = await createProjectResolver(projectDir);
+  const normalizedName = normalizeModuleName(moduleName);
+  const ast = resolver.resolveModule(normalizedName);
+  if (!ast) throw new Error(`Could not resolve module: ${moduleName}`);
+  return ast.srcModule.src;
+}
+
+/** Resolves shader source from either inline src or moduleName.
+ * @param src Inline shader source code
+ * @param moduleName Shader module name to load from filesystem
+ * @param projectDir Project directory for module resolution
+ * @returns Shader source code
+ */
+export async function resolveShaderSource(
+  src?: string,
+  moduleName?: string,
+  projectDir?: string,
+): Promise<string> {
+  validateSourceParams(src, moduleName);
+  return moduleName
+    ? await loadShaderSourceFromModule(moduleName, projectDir)
+    : src!;
+}

package/src/TestComputeShader.ts

@@ -0,0 +1,178 @@
+import { copyBuffer, elementStride, type WgslElementType } from "thimbleberry";
+import type { LinkParams } from "wesl";
+import { withErrorScopes } from "wesl-gpu";
+import { compileShader } from "./CompileShader.ts";
+import { resolveShaderSource } from "./ShaderModuleLoader.ts";
+
+const defaultResultSize = 4; // 4 elements
+
+export interface ComputeTestParams {
+  /** WESL/WGSL source code for the compute shader to test.
+   * Either src or moduleName must be provided, but not both. */
+  src?: string;
+
+  /** Name of shader module to load from filesystem.
+   * Supports: bare name (sum.wgsl), path (algorithms/sum.wgsl), or module path (package::algorithms::sum).
+   * Either src or moduleName must be provided, but not both. */
+  moduleName?: string;
+
+  /** Project directory for resolving shader dependencies.
+   * Allows the shader to import from npm shader libraries.
+   * Optional: defaults to searching upward from cwd for package.json or wesl.toml.
+   * Typically use `import.meta.url`. */
+  projectDir?: string;
+
+  /** GPU device for running the tests.
+   * Typically use `getGPUDevice()` from wgsl-test. */
+  device: GPUDevice;
+
+  /** Format of the result buffer. Default: "u32" */
+  resultFormat?: WgslElementType;
+
+  /** Size of result buffer in elements. Default: 4 */
+  size?: number;
+
+  /** Flags for conditional compilation to test shader specialization.
+   * Useful for testing `@if` statements in the shader. */
+  conditions?: LinkParams["conditions"];
+
+  /** Constants for shader compilation.
+   * Injects host-provided values via the `constants::` namespace. */
+  constants?: LinkParams["constants"];
+
+  /** Use source shaders from current package instead of built bundles.
+   * Default: true for faster iteration during development. */
+  useSourceShaders?: boolean;
+
+  /** Number of workgroups to dispatch. Default: 1
+   * Can be a single number or [x, y, z] for multi-dimensional dispatch. */
+  dispatchWorkgroups?: number | [number, number, number];
+}
+
+/**
+ * Compiles and runs a compute shader on the GPU for testing.
+ *
+ * Provides a storage buffer available at `test::results` where the shader
+ * can write test output. After execution, the storage buffer is copied back
+ * to the CPU and returned for validation.
+ *
+ * Shader libraries mentioned in the source are automatically resolved from node_modules.
+ *
+ * @returns Array of numbers from the storage buffer (typically 4 elements for u32/f32 format)
+ */
+export async function testCompute(
+  params: ComputeTestParams,
+): Promise<number[]> {
+  const {
+    projectDir,
+    device,
+    src,
+    moduleName,
+    conditions = {},
+    constants,
+    useSourceShaders,
+    dispatchWorkgroups = 1,
+  } = params;
+  const { resultFormat = "u32", size = defaultResultSize } = params;
+
+  // Resolve shader source from either src or moduleName
+  const shaderSrc = await resolveShaderSource(src, moduleName, projectDir);
+
+  const arrayType = `array<${resultFormat}, ${size}>`;
+  const virtualLibs = {
+    test: () =>
+      `@group(0) @binding(0) var <storage, read_write> results: ${arrayType};`,
+  };
+
+  const module = await compileShader({
+    projectDir,
+    device,
+    src: shaderSrc,
+    conditions,
+    constants,
+    virtualLibs,
+    useSourceShaders,
+  });
+
+  return await runCompute(
+    device,
+    module,
+    resultFormat,
+    size,
+    dispatchWorkgroups,
+  );
+}
+
+/**
+ * Runs a compiled compute shader and returns the result buffer.
+ *
+ * Creates a storage buffer at @group(0) @binding(0) where the shader can
+ * write output. The shader is invoked once, then the buffer is copied back
+ * to the CPU for reading.
+ *
+ * @param module - The compiled GPUShaderModule containing the compute shader
+ * @param resultFormat - Format for interpreting result buffer data (default: u32)
+ * @param size - Size of result buffer in bytes (default: 16)
+ * @param dispatchWorkgroups - Number of workgroups to dispatch (default: 1)
+ * @returns Array containing the shader's output from the storage buffer
+ */
+export async function runCompute(
+  device: GPUDevice,
+  module: GPUShaderModule,
+  resultFormat: WgslElementType = "u32",
+  size = defaultResultSize,
+  dispatchWorkgroups: number | [number, number, number] = 1,
+): Promise<number[]> {
+  return await withErrorScopes(device, async () => {
+    const bgLayout = device.createBindGroupLayout({
+      entries: [
+        {
+          binding: 0,
+          visibility: GPUShaderStage.COMPUTE,
+          buffer: { type: "storage" },
+        },
+      ],
+    });
+
+    const pipeline = device.createComputePipeline({
+      layout: device.createPipelineLayout({ bindGroupLayouts: [bgLayout] }),
+      compute: { module },
+    });
+
+    const storageBuffer = createStorageBuffer(
+      device,
+      size * elementStride(resultFormat),
+    );
+    const bindGroup = device.createBindGroup({
+      layout: bgLayout,
+      entries: [{ binding: 0, resource: { buffer: storageBuffer } }],
+    });
+
+    const commands = device.createCommandEncoder();
+    const pass = commands.beginComputePass();
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, bindGroup);
+    if (typeof dispatchWorkgroups === "number") {
+      pass.dispatchWorkgroups(dispatchWorkgroups);
+    } else {
+      pass.dispatchWorkgroups(...dispatchWorkgroups);
+    }
+    pass.end();
+    device.queue.submit([commands.finish()]);
+
+    return await copyBuffer(device, storageBuffer, resultFormat);
+  });
+}
+
+function createStorageBuffer(device: GPUDevice, targetSize: number): GPUBuffer {
+  const buffer = device.createBuffer({
+    label: "storage",
+    size: targetSize,
+    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
+    mappedAtCreation: true,
+  });
+  // Sentinel values to detect unwritten results
+  new Float32Array(buffer.getMappedRange()).fill(-999.0);
+  buffer.unmap();
+  return buffer;
+}