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/packages/typegpu-testing-utility/src/capture.ts b/packages/typegpu-testing-utility/src/capture.ts new file mode 100644 index 0000000000..b2e06962b1 --- /dev/null +++ b/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: (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, + }; +} diff --git a/packages/typegpu-testing-utility/src/index.ts b/packages/typegpu-testing-utility/src/index.ts index 84a4d236fd..c07a2cbbf4 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, simplifyType } from './capture.ts'; diff --git a/packages/typegpu/src/internal.ts b/packages/typegpu/src/internal.ts index b1e0fe93f7..4e6ade4277 100644 --- a/packages/typegpu/src/internal.ts +++ b/packages/typegpu/src/internal.ts @@ -7,6 +7,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'; diff --git a/packages/typegpu/tests/internal/capturedSnippets.test.ts b/packages/typegpu/tests/internal/capturedSnippets.test.ts new file mode 100644 index 0000000000..fe06be7c3a --- /dev/null +++ b/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); + }); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 26295c0a2e..c09cedccbe 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -950,6 +950,9 @@ importers: packages/typegpu-testing-utility: dependencies: + tinyest: + specifier: workspace:* + version: link:../tinyest typegpu: specifier: workspace:* version: link:../typegpu