feat(@typegpu/prune): Precompile TypeGPU code at build time

packages/typegpu-prune/README.md

@@ -0,0 +1,56 @@
+# @typegpu/prune
+
+> Work in progress.
+
+Write TypeGPU as usual, mark the parts to compile at build time with `precompile()`, and ship
+WGSL, layouts and generated buffer codecs instead of TypeGPU.
+
+```ts
+import { tgpu, d } from 'typegpu';
+import { init, precompile } from '@typegpu/prune';
+
+const sim = precompile((root) => {
+  const Params = d.struct({ speed: d.f32 });
+  const params = root.createUniform(Params, { speed: 1 });
+  // ... shaders, layouts
+  // Targets without a format use the preferred canvas format, picked at runtime
+  const render = root.createRenderPipeline({ /* ... */ });
+  return { Params, params, render };
+});
+
+const root = await init();
+const gpu = sim(root);
+gpu.params.write({ speed: 2 });
+```
+
+- **In dev** (or without the plugin), the arrow runs at runtime on full TypeGPU.
+- **In a precompiling build**, the plugin evaluates the arrow in Node, replaces the call with the
+  result, and points `@typegpu/prune` imports at `@typegpu/prune/runtime`.
+
+## Rules for the arrow
+
+- It may use imports, and module-level constants with literal values.
+- Schemas, layouts and shader functions are defined inside of it, or imported.
+- It receives only the root. Settings known only at runtime are precompiled as variants and
+  picked from at runtime. Variants share their WGSL.
+- The canvas format is the exception: render pipeline targets without a `format` use
+  `navigator.gpu.getPreferredCanvasFormat()` at runtime, like in TypeGPU. Calling it inside of the
+  arrow is an error.
+- It may return plain data, schemas, buffers, samplers, bind group layouts, bind groups, vertex
+  layouts and compute, guarded compute, and render pipelines.
+- `console.log` in shaders is compiled out.
+
+## Setup
+
+```ts
+// vite.config.ts
+import prune from '@typegpu/prune/vite';
+
+export default { plugins: [prune()] };
+```
+
+The prune plugin runs unplugin-typegpu itself, so it replaces it in the config (pass its options
+through `prune({ typegpu: { ... } })`).
+
+Also available for Rollup, Rolldown, esbuild and webpack. Precompiling is on for `vite build`, and
+otherwise when `NODE_ENV` is `production`; pass `{ precompile: true | false }` to override.

packages/typegpu-prune/package.json

@@ -0,0 +1,53 @@
+{
+  "name": "@typegpu/prune",
+  "version": "0.0.0",
+  "private": true,
+  "description": "Precompile TypeGPU code at build time, so that TypeGPU itself never reaches the bundle.",
+  "keywords": [
+    "precompile",
+    "typegpu",
+    "webgpu",
+    "wgsl"
+  ],
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/software-mansion/TypeGPU.git",
+    "directory": "packages/typegpu-prune"
+  },
+  "type": "module",
+  "sideEffects": false,
+  "exports": {
+    "./package.json": "./package.json",
+    ".": "./src/index.ts",
+    "./runtime": "./src/runtime/index.ts",
+    "./vite": "./src/plugin/vite.ts",
+    "./rollup": "./src/plugin/rollup.ts",
+    "./rolldown": "./src/plugin/rolldown.ts",
+    "./esbuild": "./src/plugin/esbuild.ts",
+    "./webpack": "./src/plugin/webpack.ts"
+  },
+  "scripts": {
+    "test:types": "pnpm tsc --p ./tsconfig.json --noEmit"
+  },
+  "dependencies": {
+    "@babel/parser": "^7.29.0",
+    "@babel/traverse": "^7.29.0",
+    "@babel/types": "catalog:",
+    "@jridgewell/remapping": "^2.3.5",
+    "esbuild": "^0.25.11",
+    "magic-string": "^0.30.21",
+    "unplugin": "^3.0.0",
+    "unplugin-typegpu": "workspace:*"
+  },
+  "devDependencies": {
+    "@types/babel__traverse": "^7.20.7",
+    "@types/node": "catalog:types",
+    "@webgpu/types": "catalog:types",
+    "typegpu": "workspace:*",
+    "typescript": "catalog:types"
+  },
+  "peerDependencies": {
+    "typegpu": "workspace:^"
+  }
+}

packages/typegpu-prune/src/artifact.ts

@@ -0,0 +1,118 @@
+/**
+ * The shape of what a `precompile()` call turns into at build time.
+ * Shared by the build-time serializer and the pruned runtime.
+ */
+
+export type Writer = (
+  output: DataView,
+  offset: number,
+  value: unknown,
+  littleEndian?: boolean,
+  endOffset?: number,
+) => void;
+
+export type Reader = (input: DataView, offset: number, littleEndian?: boolean) => unknown;
+
+/** A schema, reduced to what the CPU needs to move its values in and out of buffers */
+export interface Codec {
+  readonly size: number;
+  readonly align: number;
+  readonly write?: Writer | undefined;
+  readonly read?: Reader | undefined;
+}
+
+export type ResourceRef =
+  | { readonly buffer: number; readonly offset?: number; readonly size?: number }
+  | { readonly sampler: number };
+
+export interface BufferArtifact {
+  readonly size: number;
+  readonly usage: number;
+  /** Index into `codecs`, when the schema of the buffer is known */
+  readonly codec?: number | undefined;
+  /** Initial contents, base64-encoded. Omitted when all zeros */
+  readonly bytes?: string | undefined;
+}
+
+export interface LayoutArtifact {
+  /** Key of each entry, in binding order (`null` for skipped bindings) */
+  readonly keys: (string | null)[];
+  readonly entries: GPUBindGroupLayoutEntry[];
+}
+
+export interface BindGroupArtifact {
+  readonly layout: number;
+  readonly entries: (ResourceRef | null)[];
+}
+
+export interface VertexLayoutArtifact {
+  readonly arrayStride: number;
+  readonly stepMode: GPUVertexStepMode;
+}
+
+/** A bind group slot of a pipeline: which layout it expects, and a bind group if one is fixed */
+export interface GroupArtifact {
+  readonly layout: number;
+  readonly bindGroup?: number | undefined;
+}
+
+export interface ComputePipelineArtifact {
+  readonly shader: number;
+  readonly entryPoint?: string | undefined;
+  readonly groups: (GroupArtifact | null)[];
+}
+
+export interface GuardedPipelineArtifact {
+  readonly pipeline: number;
+  /** Buffer holding the `vec3u` size uniform */
+  readonly sizeBuffer: number;
+  readonly workgroupSize: [number, number, number];
+}
+
+export interface RenderPipelineArtifact {
+  readonly shader: number;
+  readonly groups: (GroupArtifact | null)[];
+  /** Vertex layout of each vertex buffer slot */
+  readonly vertexSlots: number[];
+  /** Vertex buffers bound at build time, as [vertex layout, buffer] */
+  readonly vertexBuffers: [number, number][];
+  /** The pipeline descriptor without `layout` and shader modules */
+  readonly descriptor: {
+    readonly vertex: Omit<GPUVertexState, 'module'>;
+    readonly fragment?: Omit<GPUFragmentState, 'module'>;
+    readonly primitive?: GPUPrimitiveState;
+    readonly depthStencil?: GPUDepthStencilState;
+    readonly multisample?: GPUMultisampleState;
+  };
+  /**
+   * Fragment targets that were left without a format. Their format is the preferred
+   * canvas format, known only at runtime.
+   */
+  readonly preferredFormatTargets?: number[] | undefined;
+}
+
+export interface Tables {
+  readonly codecs: Codec[];
+  readonly shaders: string[];
+  readonly buffers: BufferArtifact[];
+  readonly samplers: GPUSamplerDescriptor[];
+  readonly layouts: LayoutArtifact[];
+  readonly bindGroups: BindGroupArtifact[];
+  readonly vertexLayouts: VertexLayoutArtifact[];
+  readonly computePipelines: ComputePipelineArtifact[];
+  readonly guardedPipelines: GuardedPipelineArtifact[];
+  readonly renderPipelines: RenderPipelineArtifact[];
+}
+
+/** Looks up runtime objects for table entries, used by the generated code that rebuilds the returned value */
+export interface AttachContext {
+  codec(id: number): Codec;
+  buffer(id: number): unknown;
+  sampler(id: number): unknown;
+  layout(id: number): unknown;
+  bindGroup(id: number): unknown;
+  vertexLayout(id: number): unknown;
+  compute(id: number): unknown;
+  guarded(id: number): unknown;
+  render(id: number): unknown;
+}

packages/typegpu-prune/src/compiler/evaluate.ts

@@ -0,0 +1,122 @@
+/**
+ * The build-time evaluator: bundles a `precompile()` arrow together with the serializer
+ * for Node (applying unplugin-typegpu), runs it, and returns the generated code.
+ */
+
+import { build, type Plugin } from 'esbuild';
+import { createHash } from 'node:crypto';
+import { existsSync } from 'node:fs';
+import { mkdir, realpath, rm, writeFile } from 'node:fs/promises';
+import { tmpdir } from 'node:os';
+import { dirname, join } from 'node:path';
+import { fileURLToPath, pathToFileURL } from 'node:url';
+import typegpuPlugin from 'unplugin-typegpu/esbuild';
+import type { SerializeResult } from './serialize.ts';
+import type { Options as TypegpuOptions } from 'unplugin-typegpu';
+
+function serializerPath(): string {
+  // Source in the monorepo, compiled output once published
+  const candidates = ['./serialize.ts', './serialize.js', './serialize.mjs'].map((file) =>
+    fileURLToPath(new URL(file, import.meta.url)),
+  );
+  const found = candidates.find((file) => existsSync(file));
+  if (!found) {
+    throw new Error('[@typegpu/prune] Could not find the build-time serializer.');
+  }
+  return found;
+}
+
+/** Resolves imports of the virtual arrow module, and every `typegpu` import, as if they came from the user's file */
+function resolveFromUserFile(arrowFile: string, userDir: string): Plugin {
+  return {
+    name: 'typegpu-prune:resolve',
+    setup(ctx) {
+      ctx.onResolve({ filter: /.*/ }, async (args) => {
+        const fromArrow = args.importer === arrowFile;
+        const isTypegpu = /^typegpu(\/.*)?$/.test(args.path);
+        if ((!fromArrow && !isTypegpu) || args.pluginData?.prune) {
+          return undefined;
+        }
+        const result = await ctx.resolve(args.path, {
+          kind: args.kind,
+          resolveDir: userDir,
+          pluginData: { prune: true },
+        });
+        if (result.errors.length > 0) {
+          return { errors: result.errors };
+        }
+        return { path: result.path, external: result.external, namespace: result.namespace };
+      });
+    },
+  };
+}
+
+export interface EvaluateOptions {
+  /** Absolute path of the file the call is in */
+  filename: string;
+  /** Source of the virtual module: imports, copied constants, `export default <arrow>` and `export const options` */
+  moduleSource: string;
+  /** Options for unplugin-typegpu, applied to everything the arrow imports */
+  typegpuOptions?: TypegpuOptions | undefined;
+}
+
+export async function evaluatePrecompile(options: EvaluateOptions): Promise<SerializeResult> {
+  const hash = createHash('sha256')
+    .update(options.filename)
+    .update(options.moduleSource)
+    .digest('hex')
+    .slice(0, 16);
+  const requestedDir = join(tmpdir(), 'typegpu-prune', `${hash}-${process.pid}-${Date.now()}`);
+  await mkdir(requestedDir, { recursive: true });
+  // esbuild reports importers by their real path (the temp dir is a symlink on macOS)
+  const workDir = await realpath(requestedDir);
+
+  const ext = /\.[cm]?tsx$|\.jsx$/.test(options.filename) ? 'tsx' : 'ts';
+  const arrowFile = join(workDir, `arrow.${ext}`);
+  const entryFile = join(workDir, 'entry.ts');
+  const outFile = join(workDir, 'out.mjs');
+
+  try {
+    await writeFile(arrowFile, options.moduleSource);
+    await writeFile(
+      entryFile,
+      `import { serializePrecompiled } from ${JSON.stringify(serializerPath())};
+import fn, { options } from ${JSON.stringify(arrowFile)};
+export default () => serializePrecompiled(fn, options);
+`,
+    );
+
+    await build({
+      entryPoints: [entryFile],
+      outfile: outFile,
+      bundle: true,
+      format: 'esm',
+      platform: 'node',
+      target: 'node20',
+      sourcemap: 'inline',
+      logLevel: 'silent',
+      plugins: [
+        resolveFromUserFile(arrowFile, dirname(options.filename)),
+        typegpuPlugin(options.typegpuOptions ?? {}),
+      ],
+    });
+
+    // Every evaluation brings its own copy of TypeGPU, which would otherwise warn about duplicates
+    const globals = globalThis as { __TYPEGPU_VERSION__?: unknown };
+    const hostVersion = globals.__TYPEGPU_VERSION__;
+    delete globals.__TYPEGPU_VERSION__;
+    let mod: { default: () => Promise<SerializeResult> };
+    try {
+      mod = await import(`${pathToFileURL(outFile).href}?t=${Date.now()}`);
+    } finally {
+      if (hostVersion === undefined) {
+        delete globals.__TYPEGPU_VERSION__;
+      } else {
+        globals.__TYPEGPU_VERSION__ = hostVersion;
+      }
+    }
+    return await mod.default();
+  } finally {
+    await rm(workDir, { recursive: true, force: true });
+  }
+}