test: capture resolved statements

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

@@ -1,11 +1,21 @@
-import { UnknownData, WgslGenerator, type Snippet, dualImpl } from 'typegpu/~internal';
+import {
+  UnknownData,
+  WgslGenerator,
+  type FunctionDefinitionOptions,
+  type ResolvedStatement,
+  type Snippet,
+  dualImpl,
+} from 'typegpu/~internal';
 import * as tinyest from 'tinyest';
 import { tgpu, type TgpuFn } from 'typegpu';
+import { Void } from 'typegpu/data';
 
 const { NodeTypeCatalog: NODE } = tinyest;
 
 export class CapturingGenerator extends WgslGenerator {
   public capturedSnippets: Snippet[] = [];
+  public capturedStatements: ResolvedStatement[] = [];
+  #captureFollowingByBlock: boolean[] = [];
 
   protected _expression(expression: tinyest.Expression): Snippet {
     if (Array.isArray(expression) && expression[0] === NODE.call) {
@@ -19,6 +29,67 @@ export class CapturingGenerator extends WgslGenerator {
     }
     return super._expression(expression);
   }
+
+  protected _statement(statement: tinyest.Statement): ResolvedStatement {
+    const currentBlock = this.#captureFollowingByBlock.length - 1;
+    if (Array.isArray(statement) && statement[0] === NODE.call) {
+      const [_, calleeNode, argNodes] = statement;
+      const callee = this._expression(calleeNode);
+      if (callee.value === CAPTURE_FOLLOWING && argNodes.length === 0) {
+        if (currentBlock < 0) {
+          throw new Error('CAPTURE_FOLLOWING can only be used inside a function');
+        }
+        if (this.#captureFollowingByBlock[currentBlock]) {
+          throw new Error('CAPTURE_FOLLOWING must be followed by a statement');
+        }
+        this.#captureFollowingByBlock[currentBlock] = true;
+        return { code: '', definesInNearestScope: false };
+      }
+    }
+
+    const shouldCapture = this.#captureFollowingByBlock[currentBlock] === true;
+    if (shouldCapture) {
+      this.#captureFollowingByBlock[currentBlock] = false;
+    }
+
+    const resolved = super._statement(statement);
+    if (shouldCapture) {
+      this.capturedStatements.push(resolved);
+    }
+    return resolved;
+  }
+
+  protected _block(
+    block: tinyest.Block,
+    allowInlining: boolean,
+    externalMap?: Record<string, unknown>,
+  ): ResolvedStatement {
+    this.#captureFollowingByBlock.push(false);
+    try {
+      const resolved = super._block(block, allowInlining, externalMap);
+      const currentBlock = this.#captureFollowingByBlock.length - 1;
+      if (this.#captureFollowingByBlock[currentBlock]) {
+        throw new Error('CAPTURE_FOLLOWING must be followed by a statement');
+      }
+      return resolved;
+    } finally {
+      this.#captureFollowingByBlock.pop();
+    }
+  }
+
+  public functionDefinition(options: FunctionDefinitionOptions): string {
+    const firstCapturedStatement = this.capturedStatements.length;
+    const definition = super.functionDefinition(options);
+
+    for (let i = firstCapturedStatement; i < this.capturedStatements.length; i++) {
+      const statement = this.capturedStatements[i];
+      if (statement) {
+        statement.code = this._replaceVariablePlaceholders(statement.code);
+      }
+    }
+
+    return definition;
+  }
 }
 
 export const CAPTURE = dualImpl({
@@ -29,6 +100,14 @@ export const CAPTURE = dualImpl({
   sideEffects: false,
 });
 
+export const CAPTURE_FOLLOWING = dualImpl<() => void>({
+  name: 'CAPTURE_FOLLOWING',
+  signature: { argTypes: [], returnType: Void },
+  normalImpl: () => undefined,
+  codegenImpl: () => '',
+  sideEffects: false,
+});
+
 export function captureSnippets(fn: TgpuFn | (() => unknown)) {
   const generator = new CapturingGenerator();
 
@@ -37,6 +116,15 @@ export function captureSnippets(fn: TgpuFn | (() => unknown)) {
   return generator.capturedSnippets;
 }
 
+/** Captures the next resolved statement, including an empty statement if it folds away at comptime. */
+export function captureStatements(fn: TgpuFn | (() => unknown)) {
+  const generator = new CapturingGenerator();
+
+  tgpu.resolve([fn], { unstable_shaderGenerator: generator });
+
+  return generator.capturedStatements;
+}
+
 export function simplifyType(snippet: Snippet) {
   return {
     ...snippet,

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

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

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -1101,21 +1101,7 @@ export class WgslGenerator implements ShaderGenerator {
       `Expecting exactly ${functionInitialBlockDepth - 1} block(s) before going into the first function block scope`,
     );
     let body = this._block(options.body, /* allowInlining */ false);
-    const scope = this.ctx.topFunctionScope;
-    invariant(scope, 'Expected function scope to be present');
-    const replacements = Object.fromEntries(
-      [...scope.placeholderForVariable.entries()].map(([variable, placeholder]) => [
-        placeholder,
-        scope.modifiedVariables.has(variable) ? 'var' : 'let',
-      ]),
-    );
-    if (Object.keys(replacements).length > 0) {
-      const regex = new RegExp(Object.keys(replacements).join('|'), 'gi');
-      body.code = body.code.replace(
-        regex,
-        (match) => replacements[match as keyof typeof replacements] ?? '#ERR',
-      );
-    }
+    body.code = this._replaceVariablePlaceholders(body.code);
 
     // Only after generating the body can we determine the return type
     const returnType = options.determineReturnType();
@@ -1148,6 +1134,25 @@ export class WgslGenerator implements ShaderGenerator {
     return `${attributes}fn ${options.name}${head}${body.code || '{}'}`;
   }
 
+  protected _replaceVariablePlaceholders(code: string): string {
+    const scope = this.ctx.topFunctionScope;
+    invariant(scope, 'Expected function scope to be present');
+    const replacements = Object.fromEntries(
+      [...scope.placeholderForVariable.entries()].map(([variable, placeholder]) => [
+        placeholder,
+        scope.modifiedVariables.has(variable) ? 'var' : 'let',
+      ]),
+    );
+    if (Object.keys(replacements).length > 0) {
+      const regex = new RegExp(Object.keys(replacements).join('|'), 'gi');
+      return code.replace(
+        regex,
+        (match) => replacements[match as keyof typeof replacements] ?? '#ERR',
+      );
+    }
+    return code;
+  }
+
   /**
    * Generates a WGSL type string for the given data type, and adds necessary
    * definitions to the shader preamble. This shouldn't be called directly, only

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

@@ -1,6 +1,13 @@
 import { describe, expect } from 'vitest';
 import { tgpu, d } from 'typegpu';
-import { CAPTURE, captureSnippets, it, simplifyType } from 'typegpu-testing-utility';
+import {
+  CAPTURE,
+  CAPTURE_FOLLOWING,
+  captureSnippets,
+  captureStatements,
+  it,
+  simplifyType,
+} from 'typegpu-testing-utility';
 
 describe('CAPTURE', () => {
   it('is a no-op in regular resolves', () => {
@@ -120,3 +127,131 @@ describe('CAPTURE', () => {
     expect(captureSnippets(fn)[0]?.value).toBe(1.5);
   });
 });
+
+describe('CAPTURE_FOLLOWING', () => {
+  it('is a no-op in regular resolves', () => {
+    const withCapture = tgpu.fn(
+      [d.u32],
+      d.u32,
+    )((x) => {
+      'use gpu';
+      CAPTURE_FOLLOWING();
+      return x + 1;
+    });
+    const withoutCapture = tgpu.fn(
+      [d.u32],
+      d.u32,
+    )((x) => {
+      'use gpu';
+      return x + 1;
+    });
+
+    const normalizeName = (code: string) => code.replace(/fn \w+/, 'fn captured');
+    expect(normalizeName(tgpu.resolve([withCapture]))).toBe(
+      normalizeName(tgpu.resolve([withoutCapture])),
+    );
+  });
+
+  it('captures the following resolved statement', () => {
+    const fn = tgpu.fn(
+      [d.u32],
+      d.u32,
+    )((x) => {
+      'use gpu';
+      CAPTURE_FOLLOWING();
+      const y = x + 1;
+      return y;
+    });
+
+    expect(captureStatements(fn)).toEqual([
+      {
+        code: '  let y = (x + 1u);',
+        definesInNearestScope: true,
+      },
+    ]);
+  });
+
+  it('captures an outer statement instead of its nested statements', () => {
+    const fn = tgpu.fn(
+      [d.u32],
+      d.u32,
+    )((x) => {
+      'use gpu';
+      CAPTURE_FOLLOWING();
+      if (x > 0) {
+        return x;
+      }
+      return 0;
+    });
+
+    const captured = captureStatements(fn);
+    expect(captured).toHaveLength(1);
+    expect(captured[0]?.code).toContain('if (');
+    expect(captured[0]?.code).toContain('return x;');
+  });
+
+  it('rejects a marker without a following statement', () => {
+    const fn = () => {
+      'use gpu';
+      CAPTURE_FOLLOWING();
+    };
+
+    expect(() => captureStatements(fn)).toThrow(
+      'CAPTURE_FOLLOWING must be followed by a statement',
+    );
+  });
+
+  it('does not carry an unfinished capture into another function', () => {
+    const unfinished = () => {
+      'use gpu';
+      CAPTURE_FOLLOWING();
+    };
+    const unrelated = () => {
+      'use gpu';
+      const value = 1;
+    };
+    const fn = () => {
+      'use gpu';
+      unfinished();
+      unrelated();
+    };
+
+    expect(() => captureStatements(fn)).toThrow(
+      'CAPTURE_FOLLOWING must be followed by a statement',
+    );
+  });
+
+  it('rejects consecutive markers', () => {
+    const fn = () => {
+      'use gpu';
+      CAPTURE_FOLLOWING();
+      CAPTURE_FOLLOWING();
+      const value = 1;
+    };
+
+    expect(() => captureStatements(fn)).toThrow(
+      'CAPTURE_FOLLOWING must be followed by a statement',
+    );
+  });
+
+  it('resolves deferred variable placeholders in captured code', () => {
+    const fn = tgpu.fn([d.vec3f])((x) => {
+      'use gpu';
+      CAPTURE_FOLLOWING();
+      const vector = x + d.vec3f(1);
+      CAPTURE_FOLLOWING();
+      for (let i = 0; i < 2; i++) {
+        vector + d.vec3f(i);
+      }
+    });
+
+    const captured = captureStatements(fn);
+    const shader = tgpu.resolve([fn]);
+
+    expect(captured).toHaveLength(2);
+    for (const statement of captured) {
+      expect(statement.code).not.toContain('#VAR_');
+      expect(shader).toContain(statement.code.trim());
+    }
+  });
+});