From bcf1dfbc9c4a668c5513db92d8541abfa519c5be Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Mon, 13 Jul 2026 13:40:52 +0200 Subject: [PATCH 1/9] Add capturing stubs --- .../typegpu-testing-utility/src/capture.ts | 32 +++++++++++++++++++ packages/typegpu-testing-utility/src/index.ts | 1 + 2 files changed, 33 insertions(+) create mode 100644 packages/typegpu-testing-utility/src/capture.ts diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts new file mode 100644 index 0000000000..4bb970cbdb --- /dev/null +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -0,0 +1,32 @@ +import { WgslGenerator, type Snippet } from 'typegpu/~internal'; +import type { Expression } from '../../tinyest/src/nodes'; +import * as tinyest from 'tinyest'; +import { tgpu } from 'typegpu'; + +const { NodeTypeCatalog: NODE } = tinyest; + +export class CapturingGenerator extends WgslGenerator { + // captured snippets + + protected _expression(expression: Expression): Snippet { + if (Array.isArray(expression) && expression[0] === NODE.call) { + const [_, calleeNode, argNodes] = expression; + const callee = this._expression(calleeNode); + if (callee.value === CAPTURE) { + console.log('CAPTURING'); + } + return super._expression(argNodes[0]); + } + return super._expression(expression); + } +} + +export const CAPTURE = (a: T): T => a; + +export function captureSnippets(fn: () => unknown) { + const generator = new CapturingGenerator(); + + tgpu.resolve([fn], { unstable_shaderGenerator: generator }); + + return; +} diff --git a/packages/typegpu-testing-utility/src/index.ts b/packages/typegpu-testing-utility/src/index.ts index 84a4d236fd..2a6426bdff 100644 --- a/packages/typegpu-testing-utility/src/index.ts +++ b/packages/typegpu-testing-utility/src/index.ts @@ -1 +1,2 @@ export { it, test } from './extendedIt.ts'; +export { CAPTURE, captureSnippets } from './capture.ts'; From 237a13b33ff8abf5a4fc87907190db8e65a53602 Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Mon, 13 Jul 2026 13:47:11 +0200 Subject: [PATCH 2/9] Allow for resolving capture --- .../typegpu-testing-utility/src/capture.ts | 10 ++++++- .../tests/internal/capturedSnippets.test.ts | 27 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 packages/typegpu/tests/internal/capturedSnippets.test.ts diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts index 4bb970cbdb..31f8d9009d 100644 --- a/packages/typegpu-testing-utility/src/capture.ts +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -2,6 +2,8 @@ import { WgslGenerator, type Snippet } from 'typegpu/~internal'; import type { Expression } from '../../tinyest/src/nodes'; import * as tinyest from 'tinyest'; import { tgpu } from 'typegpu'; +import { dualImpl } from '../../typegpu/src/core/function/dualImpl.ts'; +import { stitch } from '../../typegpu/src/core/resolve/stitch.ts'; const { NodeTypeCatalog: NODE } = tinyest; @@ -21,7 +23,13 @@ export class CapturingGenerator extends WgslGenerator { } } -export const CAPTURE = (a: T): T => a; +export const CAPTURE = dualImpl({ + name: 'CAPTURE', + signature: (arg) => ({ argTypes: [arg], returnType: arg }), + normalImpl: (expr: T): T => expr, + codegenImpl: (_ctx, [expr]) => stitch`${expr}`, + sideEffects: false, +}); export function captureSnippets(fn: () => unknown) { const generator = new CapturingGenerator(); diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts new file mode 100644 index 0000000000..3cffb9f336 --- /dev/null +++ b/packages/typegpu/tests/internal/capturedSnippets.test.ts @@ -0,0 +1,27 @@ +// TODO(#2659): Move out of /internal once `getName` is available through 'typegpu/~internal' +import { describe, expect } from 'vitest'; +import { struct } from 'typegpu/data'; +import { getName } from '../../src/shared/meta.ts'; +import { tgpu, d, type TgpuBindGroupLayout } from 'typegpu'; +import { CAPTURE, captureSnippets, it } from 'typegpu-testing-utility'; + +describe('...', () => { + 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; + }" + `); + }); +}); From 12faa7e030db7efff520d129e3ca02d059b888d7 Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:02:12 +0200 Subject: [PATCH 3/9] Allow snippet capturing --- .../typegpu-testing-utility/src/capture.ts | 13 ++-- .../tests/internal/capturedSnippets.test.ts | 68 +++++++++++++++++++ 2 files changed, 75 insertions(+), 6 deletions(-) diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts index 31f8d9009d..7063dec781 100644 --- a/packages/typegpu-testing-utility/src/capture.ts +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -1,23 +1,24 @@ import { WgslGenerator, type Snippet } from 'typegpu/~internal'; import type { Expression } from '../../tinyest/src/nodes'; import * as tinyest from 'tinyest'; -import { tgpu } from 'typegpu'; +import { tgpu, type TgpuFn } from 'typegpu'; import { dualImpl } from '../../typegpu/src/core/function/dualImpl.ts'; import { stitch } from '../../typegpu/src/core/resolve/stitch.ts'; const { NodeTypeCatalog: NODE } = tinyest; export class CapturingGenerator extends WgslGenerator { - // captured snippets + public capturedSnippets: Snippet[] = []; protected _expression(expression: Expression): Snippet { if (Array.isArray(expression) && expression[0] === NODE.call) { const [_, calleeNode, argNodes] = expression; const callee = this._expression(calleeNode); if (callee.value === CAPTURE) { - console.log('CAPTURING'); + const snippet = this._expression(argNodes[0]); + this.capturedSnippets.push(snippet); + return snippet; } - return super._expression(argNodes[0]); } return super._expression(expression); } @@ -31,10 +32,10 @@ export const CAPTURE = dualImpl({ sideEffects: false, }); -export function captureSnippets(fn: () => unknown) { +export function captureSnippets(fn: TgpuFn | (() => unknown)) { const generator = new CapturingGenerator(); tgpu.resolve([fn], { unstable_shaderGenerator: generator }); - return; + return generator.capturedSnippets; } diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts index 3cffb9f336..ae7c1beaff 100644 --- a/packages/typegpu/tests/internal/capturedSnippets.test.ts +++ b/packages/typegpu/tests/internal/capturedSnippets.test.ts @@ -24,4 +24,72 @@ describe('...', () => { }" `); }); + + 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)).toMatchInlineSnapshot(` + [ + SnippetImpl { + "dataType": { + "concretized": [Function], + "toString": [Function], + "type": "abstractInt", + Symbol(typegpu:0.11.9:$internal): {}, + }, + "origin": "constant", + "possibleSideEffects": true, + "value": 3, + }, + SnippetImpl { + "dataType": [Function], + "origin": "runtime", + "possibleSideEffects": false, + "value": "(a + 1i)", + }, + SnippetImpl { + "dataType": [Function], + "origin": "argument", + "possibleSideEffects": false, + "value": "x", + }, + SnippetImpl { + "dataType": { + "concretized": [Function], + "toString": [Function], + "type": "abstractInt", + Symbol(typegpu:0.11.9:$internal): {}, + }, + "origin": "constant", + "possibleSideEffects": false, + "value": 1, + }, + SnippetImpl { + "dataType": [Function], + "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); + }); }); From bb8b46e35cb3ec44803667c3f390e4aa9163f6f1 Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Mon, 13 Jul 2026 15:43:51 +0200 Subject: [PATCH 4/9] Export dualImpl from internal endpoint --- packages/typegpu-testing-utility/src/capture.ts | 8 +++----- packages/typegpu/src/internal.ts | 1 + 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts index 7063dec781..5ef7616fe8 100644 --- a/packages/typegpu-testing-utility/src/capture.ts +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -1,16 +1,14 @@ import { WgslGenerator, type Snippet } from 'typegpu/~internal'; -import type { Expression } from '../../tinyest/src/nodes'; import * as tinyest from 'tinyest'; import { tgpu, type TgpuFn } from 'typegpu'; -import { dualImpl } from '../../typegpu/src/core/function/dualImpl.ts'; -import { stitch } from '../../typegpu/src/core/resolve/stitch.ts'; +import { dualImpl } from 'typegpu/~internal'; const { NodeTypeCatalog: NODE } = tinyest; export class CapturingGenerator extends WgslGenerator { public capturedSnippets: Snippet[] = []; - protected _expression(expression: Expression): Snippet { + protected _expression(expression: tinyest.Expression): Snippet { if (Array.isArray(expression) && expression[0] === NODE.call) { const [_, calleeNode, argNodes] = expression; const callee = this._expression(calleeNode); @@ -28,7 +26,7 @@ export const CAPTURE = dualImpl({ name: 'CAPTURE', signature: (arg) => ({ argTypes: [arg], returnType: arg }), normalImpl: (expr: T): T => expr, - codegenImpl: (_ctx, [expr]) => stitch`${expr}`, + codegenImpl: (ctx, [expr]) => ctx.resolveSnippet(expr).value, sideEffects: false, }); diff --git a/packages/typegpu/src/internal.ts b/packages/typegpu/src/internal.ts index 348233398a..89af271d75 100644 --- a/packages/typegpu/src/internal.ts +++ b/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'; From 14d81aff6bf566bba730d7191ba7cc46fc2e43ac Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:01:26 +0200 Subject: [PATCH 5/9] Add `simplifyType` --- .../typegpu-testing-utility/src/capture.ts | 9 +++- packages/typegpu-testing-utility/src/index.ts | 2 +- .../tests/internal/capturedSnippets.test.ts | 52 +++++++++---------- 3 files changed, 34 insertions(+), 29 deletions(-) diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts index 5ef7616fe8..243edcde27 100644 --- a/packages/typegpu-testing-utility/src/capture.ts +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -1,4 +1,4 @@ -import { WgslGenerator, type Snippet } from 'typegpu/~internal'; +import { UnknownData, WgslGenerator, type Snippet } from 'typegpu/~internal'; import * as tinyest from 'tinyest'; import { tgpu, type TgpuFn } from 'typegpu'; import { dualImpl } from 'typegpu/~internal'; @@ -37,3 +37,10 @@ export function captureSnippets(fn: TgpuFn | (() => unknown)) { return generator.capturedSnippets; } + +export function simplifyType(snippet: Snippet) { + return { + ...snippet, + dataType: snippet.dataType === UnknownData ? 'UnknownData' : snippet.dataType.type, + }; +} diff --git a/packages/typegpu-testing-utility/src/index.ts b/packages/typegpu-testing-utility/src/index.ts index 2a6426bdff..c07a2cbbf4 100644 --- a/packages/typegpu-testing-utility/src/index.ts +++ b/packages/typegpu-testing-utility/src/index.ts @@ -1,2 +1,2 @@ export { it, test } from './extendedIt.ts'; -export { CAPTURE, captureSnippets } from './capture.ts'; +export { CAPTURE, captureSnippets, simplifyType } from './capture.ts'; diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts index ae7c1beaff..c58faf87fa 100644 --- a/packages/typegpu/tests/internal/capturedSnippets.test.ts +++ b/packages/typegpu/tests/internal/capturedSnippets.test.ts @@ -1,11 +1,8 @@ -// TODO(#2659): Move out of /internal once `getName` is available through 'typegpu/~internal' import { describe, expect } from 'vitest'; -import { struct } from 'typegpu/data'; -import { getName } from '../../src/shared/meta.ts'; -import { tgpu, d, type TgpuBindGroupLayout } from 'typegpu'; -import { CAPTURE, captureSnippets, it } from 'typegpu-testing-utility'; +import { tgpu, d } from 'typegpu'; +import { CAPTURE, captureSnippets, it, simplifyType } from 'typegpu-testing-utility'; -describe('...', () => { +describe('CAPTURE', () => { it('is a no-op in regular resolves', () => { const fn = tgpu.fn([d.u32])((x) => { 'use gpu'; @@ -34,44 +31,34 @@ describe('...', () => { const d = CAPTURE(CAPTURE(1) + (c + x)); }); - expect(captureSnippets(fn)).toMatchInlineSnapshot(` + expect(captureSnippets(fn).map(simplifyType)).toMatchInlineSnapshot(` [ - SnippetImpl { - "dataType": { - "concretized": [Function], - "toString": [Function], - "type": "abstractInt", - Symbol(typegpu:0.11.9:$internal): {}, - }, + { + "dataType": "abstractInt", "origin": "constant", "possibleSideEffects": true, "value": 3, }, - SnippetImpl { - "dataType": [Function], + { + "dataType": "i32", "origin": "runtime", "possibleSideEffects": false, "value": "(a + 1i)", }, - SnippetImpl { - "dataType": [Function], + { + "dataType": "u32", "origin": "argument", "possibleSideEffects": false, "value": "x", }, - SnippetImpl { - "dataType": { - "concretized": [Function], - "toString": [Function], - "type": "abstractInt", - Symbol(typegpu:0.11.9:$internal): {}, - }, + { + "dataType": "abstractInt", "origin": "constant", "possibleSideEffects": false, "value": 1, }, - SnippetImpl { - "dataType": [Function], + { + "dataType": "u32", "origin": "runtime", "possibleSideEffects": false, "value": "(1u + (c + x))", @@ -92,4 +79,15 @@ describe('...', () => { 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); + }); }); From 8fc5aa4d435ebd017cbbcfb28ab09d08c223a1dd Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:59:54 +0200 Subject: [PATCH 6/9] Fix package.json --- packages/typegpu-testing-utility/package.json | 1 + pnpm-lock.yaml | 3 +++ 2 files changed, 4 insertions(+) diff --git a/packages/typegpu-testing-utility/package.json b/packages/typegpu-testing-utility/package.json index 73b5d9d994..681710e6a2 100644 --- a/packages/typegpu-testing-utility/package.json +++ b/packages/typegpu-testing-utility/package.json @@ -12,6 +12,7 @@ "test:types": "pnpm tsc --p ./tsconfig.json --noEmit" }, "dependencies": { + "tinyest": "workspace:*", "typegpu": "workspace:*" }, "devDependencies": { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 1f2fb169d7..f40cb11f0a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -937,6 +937,9 @@ importers: packages/typegpu-testing-utility: dependencies: + tinyest: + specifier: workspace:* + version: link:../tinyest typegpu: specifier: workspace:* version: link:../typegpu From 00a5ba86062d862dd1dd211e009540e3f41aa56c Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:24:40 +0200 Subject: [PATCH 7/9] Add a test for _typedExpression casting --- .../typegpu/tests/internal/capturedSnippets.test.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts index c58faf87fa..9d1359e53d 100644 --- a/packages/typegpu/tests/internal/capturedSnippets.test.ts +++ b/packages/typegpu/tests/internal/capturedSnippets.test.ts @@ -90,4 +90,16 @@ describe('CAPTURE', () => { expect(captured[0]?.value).toBe(1); expect(captured[1]?.value).toBe(3); }); + + 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); + }); }); From bc3fb3ef84762d85b9f3c7e95f4a03ad0eee781e Mon Sep 17 00:00:00 2001 From: Aleksander Katan <56294622+aleksanderkatan@users.noreply.github.com> Date: Tue, 14 Jul 2026 17:38:58 +0200 Subject: [PATCH 8/9] Review fixes --- packages/typegpu-testing-utility/src/capture.ts | 3 +-- .../tests/internal/capturedSnippets.test.ts | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts index 243edcde27..b2e06962b1 100644 --- a/packages/typegpu-testing-utility/src/capture.ts +++ b/packages/typegpu-testing-utility/src/capture.ts @@ -1,7 +1,6 @@ -import { UnknownData, WgslGenerator, type Snippet } from 'typegpu/~internal'; +import { UnknownData, WgslGenerator, type Snippet, dualImpl } from 'typegpu/~internal'; import * as tinyest from 'tinyest'; import { tgpu, type TgpuFn } from 'typegpu'; -import { dualImpl } from 'typegpu/~internal'; const { NodeTypeCatalog: NODE } = tinyest; diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts index 9d1359e53d..997837f1c1 100644 --- a/packages/typegpu/tests/internal/capturedSnippets.test.ts +++ b/packages/typegpu/tests/internal/capturedSnippets.test.ts @@ -91,6 +91,23 @@ describe('CAPTURE', () => { 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( [], From 8e9cea0fd8f1318b0e1f6e4d5c3a74609f3fa23c Mon Sep 17 00:00:00 2001 From: "pullfrog[bot]" <226033991+pullfrog[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 10:01:31 +0000 Subject: [PATCH 9/9] test: Update stale snapshot after merge with main --- packages/typegpu/tests/internal/capturedSnippets.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts index 997837f1c1..fe06be7c3a 100644 --- a/packages/typegpu/tests/internal/capturedSnippets.test.ts +++ b/packages/typegpu/tests/internal/capturedSnippets.test.ts @@ -36,7 +36,7 @@ describe('CAPTURE', () => { { "dataType": "abstractInt", "origin": "constant", - "possibleSideEffects": true, + "possibleSideEffects": false, "value": 3, }, {