From dc530ca5d8060722db25a70d93f8496a07f9b3bf Mon Sep 17 00:00:00 2001 From: huymobile Date: Sun, 30 Aug 2026 14:52:21 +0700 Subject: [PATCH 1/2] test: capture resolved statements --- .../typegpu-testing-utility/src/capture.ts | 66 ++++++++++- packages/typegpu-testing-utility/src/index.ts | 8 +- .../tests/internal/capturedSnippets.test.ts | 103 +++++++++++++++++- 3 files changed, 174 insertions(+), 3 deletions(-) diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts index b2e06962b1..c6da34f15a 100644 --- a/packages/typegpu-testing-utility/src/capture.ts +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -1,11 +1,20 @@ -import { UnknownData, WgslGenerator, type Snippet, dualImpl } from 'typegpu/~internal'; +import { + UnknownData, + WgslGenerator, + 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) { @@ -16,9 +25,48 @@ export class CapturingGenerator extends WgslGenerator { this.capturedSnippets.push(snippet); return snippet; } + if (callee.value === CAPTURE_FOLLOWING && argNodes.length === 0) { + const currentBlock = this.#captureFollowingByBlock.length - 1; + if (currentBlock < 0) { + throw new Error('CAPTURE_FOLLOWING can only be used inside a function'); + } + this.#captureFollowingByBlock[currentBlock] = true; + } } return super._expression(expression); } + + protected _statement(statement: tinyest.Statement): ResolvedStatement { + const currentBlock = this.#captureFollowingByBlock.length - 1; + 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, + ): 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(); + } + } } export const CAPTURE = dualImpl({ @@ -29,6 +77,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 +93,14 @@ export function captureSnippets(fn: TgpuFn | (() => unknown)) { return generator.capturedSnippets; } +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, diff --git a/packages/typegpu-testing-utility/src/index.ts b/packages/typegpu-testing-utility/src/index.ts index c07a2cbbf4..7086b6109c 100644 --- a/packages/typegpu-testing-utility/src/index.ts +++ b/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'; diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts index fe06be7c3a..d6a4f9f257 100644 --- a/packages/typegpu/tests/internal/capturedSnippets.test.ts +++ b/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,97 @@ 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', + ); + }); +}); From acd21101cc603b3b9da6a2ddb5a660008ba0e4c7 Mon Sep 17 00:00:00 2001 From: huymobile Date: Sun, 30 Aug 2026 15:23:31 +0700 Subject: [PATCH 2/2] fix: normalize captured statements --- .../typegpu-testing-utility/src/capture.ts | 34 +++++++++++++++--- packages/typegpu/src/tgsl/wgslGenerator.ts | 35 +++++++++++-------- .../tests/internal/capturedSnippets.test.ts | 34 ++++++++++++++++++ 3 files changed, 83 insertions(+), 20 deletions(-) diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts index c6da34f15a..cf7dd44b1d 100644 --- a/packages/typegpu-testing-utility/src/capture.ts +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -1,6 +1,7 @@ import { UnknownData, WgslGenerator, + type FunctionDefinitionOptions, type ResolvedStatement, type Snippet, dualImpl, @@ -25,19 +26,27 @@ export class CapturingGenerator extends WgslGenerator { this.capturedSnippets.push(snippet); return snippet; } + } + 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) { - const currentBlock = this.#captureFollowingByBlock.length - 1; 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 }; } } - return super._expression(expression); - } - protected _statement(statement: tinyest.Statement): ResolvedStatement { - const currentBlock = this.#captureFollowingByBlock.length - 1; const shouldCapture = this.#captureFollowingByBlock[currentBlock] === true; if (shouldCapture) { this.#captureFollowingByBlock[currentBlock] = false; @@ -67,6 +76,20 @@ export class CapturingGenerator extends WgslGenerator { 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({ @@ -93,6 +116,7 @@ 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(); diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 33d64ec786..e669d315f7 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/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 diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts index d6a4f9f257..81e5ae9279 100644 --- a/packages/typegpu/tests/internal/capturedSnippets.test.ts +++ b/packages/typegpu/tests/internal/capturedSnippets.test.ts @@ -220,4 +220,38 @@ describe('CAPTURE_FOLLOWING', () => { '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()); + } + }); });