dx: More robust snippet testing

packages/typegpu-testing-utility/package.json

@@ -12,6 +12,7 @@
     "test:types": "pnpm tsc --p ./tsconfig.json --noEmit"
   },
   "dependencies": {
+    "tinyest": "workspace:*",
     "typegpu": "workspace:*"
   },
   "devDependencies": {

packages/typegpu-testing-utility/src/capture.ts

@@ -0,0 +1,45 @@
+import { UnknownData, WgslGenerator, type Snippet, dualImpl } from 'typegpu/~internal';
+import * as tinyest from 'tinyest';
+import { tgpu, type TgpuFn } from 'typegpu';
+
+const { NodeTypeCatalog: NODE } = tinyest;
+
+export class CapturingGenerator extends WgslGenerator {
+  public capturedSnippets: Snippet[] = [];
+
+  protected _expression(expression: tinyest.Expression): Snippet {
+    if (Array.isArray(expression) && expression[0] === NODE.call) {
+      const [_, calleeNode, argNodes] = expression;
+      const callee = this._expression(calleeNode);
+      if (callee.value === CAPTURE) {
+        const snippet = this._expression(argNodes[0]);
+        this.capturedSnippets.push(snippet);
+        return snippet;
+      }
+    }
+    return super._expression(expression);
+  }
+}
+
+export const CAPTURE = dualImpl({
+  name: 'CAPTURE',
+  signature: (arg) => ({ argTypes: [arg], returnType: arg }),
+  normalImpl: <T>(expr: T): T => expr,
+  codegenImpl: (ctx, [expr]) => ctx.resolveSnippet(expr).value,
+  sideEffects: false,
+});
+
+export function captureSnippets(fn: TgpuFn | (() => unknown)) {
+  const generator = new CapturingGenerator();
+
+  tgpu.resolve([fn], { unstable_shaderGenerator: generator });
+
+  return generator.capturedSnippets;
+}
+
+export function simplifyType(snippet: Snippet) {
+  return {
+    ...snippet,
+    dataType: snippet.dataType === UnknownData ? 'UnknownData' : snippet.dataType.type,
+  };
+}

packages/typegpu-testing-utility/src/index.ts

@@ -1 +1,2 @@
 export { it, test } from './extendedIt.ts';
+export { CAPTURE, captureSnippets, simplifyType } from './capture.ts';

packages/typegpu/src/internal.ts

@@ -5,6 +5,7 @@ export { UnknownData } from './data/dataTypes.ts';
 export { getName } from './shared/meta.ts';
 export { WgslGenerator } from './tgsl/wgslGenerator.ts';
 export { snip } from './data/snippet.ts';
+export { dualImpl } from './core/function/dualImpl.ts';
 
 // types
 export type { ResolutionCtx, FunctionArgument, TgpuShaderStage } from './types.ts';

packages/typegpu/tests/internal/capturedSnippets.test.ts

@@ -0,0 +1,122 @@
+import { describe, expect } from 'vitest';
+import { tgpu, d } from 'typegpu';
+import { CAPTURE, captureSnippets, it, simplifyType } from 'typegpu-testing-utility';
+
+describe('CAPTURE', () => {
+  it('is a no-op in regular resolves', () => {
+    const fn = tgpu.fn([d.u32])((x) => {
+      'use gpu';
+      const a = CAPTURE(1 + 2);
+      const b = CAPTURE(a + 1);
+      const c = CAPTURE(x);
+      const d = CAPTURE(CAPTURE(1));
+    });
+
+    expect(tgpu.resolve([fn])).toMatchInlineSnapshot(`
+      "fn fn_1(x: u32) {
+        const a = 3;
+        let b = (a + 1i);
+        let c = x;
+        const d = 1;
+      }"
+    `);
+  });
+
+  it('allows snippet extraction', () => {
+    const fn = tgpu.fn([d.u32])((x) => {
+      'use gpu';
+      const a = CAPTURE(1 + 2);
+      const b = CAPTURE(a + 1);
+      const c = CAPTURE(x);
+      const d = CAPTURE(CAPTURE(1) + (c + x));
+    });
+
+    expect(captureSnippets(fn).map(simplifyType)).toMatchInlineSnapshot(`
+      [
+        {
+          "dataType": "abstractInt",
+          "origin": "constant",
+          "possibleSideEffects": false,
+          "value": 3,
+        },
+        {
+          "dataType": "i32",
+          "origin": "runtime",
+          "possibleSideEffects": false,
+          "value": "(a + 1i)",
+        },
+        {
+          "dataType": "u32",
+          "origin": "argument",
+          "possibleSideEffects": false,
+          "value": "x",
+        },
+        {
+          "dataType": "abstractInt",
+          "origin": "constant",
+          "possibleSideEffects": false,
+          "value": 1,
+        },
+        {
+          "dataType": "u32",
+          "origin": "runtime",
+          "possibleSideEffects": false,
+          "value": "(1u + (c + x))",
+        },
+      ]
+    `);
+  });
+
+  it('recaptures when called a second time', () => {
+    let count = 0;
+    const lazy = tgpu.lazy(() => count++);
+    const fn = () => {
+      'use gpu';
+      return CAPTURE(lazy.$);
+    };
+
+    expect(captureSnippets(fn)[0]?.value).toBe(0);
+    expect(captureSnippets(fn)[0]?.value).toBe(1);
+    expect(captureSnippets(fn)[0]?.value).toBe(2);
+  });
+
+  it('captures inner to outer', () => {
+    const fn = () => {
+      'use gpu';
+      return CAPTURE(CAPTURE(1) + 2);
+    };
+
+    const captured = captureSnippets(fn);
+    expect(captured[0]?.value).toBe(1);
+    expect(captured[1]?.value).toBe(3);
+  });
+
+  it('captures structs after casting', () => {
+    const Boid = d.struct({
+      pos: d.vec3f,
+    });
+
+    const fn = tgpu.fn(
+      [],
+      Boid,
+    )(() => {
+      'use gpu';
+      return CAPTURE({ pos: d.vec3f() });
+    });
+
+    const captured = captureSnippets(fn);
+    expect(captured[0]?.dataType).toBe(Boid);
+  });
+
+  it('captures before type casting', () => {
+    const fn = tgpu.fn(
+      [],
+      d.u32,
+    )(() => {
+      'use gpu';
+      return CAPTURE(1.5);
+    });
+
+    expect(captureSnippets(fn)[0]?.value).toBe(1.5);
+  });
+});