diff --git a/packages/typegpu/src/core/function/tgpuComputeFn.ts b/packages/typegpu/src/core/function/tgpuComputeFn.ts index 89715f7c44..e505093355 100644 --- a/packages/typegpu/src/core/function/tgpuComputeFn.ts +++ b/packages/typegpu/src/core/function/tgpuComputeFn.ts @@ -27,7 +27,7 @@ type TgpuComputeFnShellHeader< readonly argTypes: [IOLayoutToSchema] | []; readonly returnType: Void; readonly workgroupSize: [number, number, number]; - readonly isEntry: true; + readonly entryPoint: 'compute'; }; /** @@ -54,22 +54,7 @@ export type TgpuComputeFnShell< & (( strings: TemplateStringsArray, ...values: unknown[] - ) => TgpuComputeFn) - & { - /** - * @deprecated Invoke the shell as a function instead. - */ - does: - & (( - implementation: (input: InferIO) => undefined, - ) => TgpuComputeFn) - & /** - * @param implementation - * Raw WGSL function implementation with header and body - * without `fn` keyword and function name - * e.g. `"(x: f32) -> f32 { return x; }"`; - */ ((implementation: string) => TgpuComputeFn); - }; + ) => TgpuComputeFn); export interface TgpuComputeFn< // biome-ignore lint/suspicious/noExplicitAny: to allow assigning any compute fn to TgpuComputeFn (non-generic) type @@ -122,7 +107,7 @@ export function computeFn< options.workgroupSize[1] ?? 1, options.workgroupSize[2] ?? 1, ], - isEntry: true, + entryPoint: 'compute', }; const call = ( @@ -135,9 +120,13 @@ export function computeFn< stripTemplate(arg, ...values), ); - return Object.assign(Object.assign(call, shell), { - does: call, - }) as TgpuComputeFnShell; + return Object.assign(call, shell); +} + +export function isTgpuComputeFn>( + value: unknown | TgpuComputeFn, +): value is TgpuComputeFn { + return (value as TgpuComputeFn)?.shell?.entryPoint === 'compute'; } // -------------- diff --git a/packages/typegpu/src/core/function/tgpuFn.ts b/packages/typegpu/src/core/function/tgpuFn.ts index cd83970ce4..f7c9b3b7a5 100644 --- a/packages/typegpu/src/core/function/tgpuFn.ts +++ b/packages/typegpu/src/core/function/tgpuFn.ts @@ -56,7 +56,6 @@ type TgpuFnShellHeader< readonly [$internal]: true; readonly argTypes: Args; readonly returnType: Return; - readonly isEntry: false; }; /** @@ -125,7 +124,6 @@ export function fn< [$internal]: true, argTypes, returnType: returnType ?? Void as Return, - isEntry: false, }; const call = ( diff --git a/packages/typegpu/src/core/function/tgpuFragmentFn.ts b/packages/typegpu/src/core/function/tgpuFragmentFn.ts index a5540a6e3e..b4cf38ad5a 100644 --- a/packages/typegpu/src/core/function/tgpuFragmentFn.ts +++ b/packages/typegpu/src/core/function/tgpuFragmentFn.ts @@ -61,7 +61,7 @@ type TgpuFragmentFnShellHeader< readonly in: FragmentIn | undefined; readonly out: FragmentOut; readonly returnType: IOLayoutToSchema; - readonly isEntry: true; + readonly entryPoint: 'fragment'; }; /** @@ -93,24 +93,7 @@ export type TgpuFragmentFnShell< & (( strings: TemplateStringsArray, ...values: unknown[] - ) => TgpuFragmentFn, FragmentOut>) - & { - /** - * @deprecated Invoke the shell as a function instead. - */ - does: - & (( - implementation: (input: InferIO) => InferIO, - ) => TgpuFragmentFn, FragmentOut>) - & /** - * @param implementation - * Raw WGSL function implementation with header and body - * without `fn` keyword and function name - * e.g. `"(x: f32) -> f32 { return x; }"`; - */ (( - implementation: string, - ) => TgpuFragmentFn, FragmentOut>); - }; + ) => TgpuFragmentFn, FragmentOut>); export interface TgpuFragmentFn< Varying extends FragmentInConstrained = FragmentInConstrained, @@ -162,7 +145,7 @@ export function fragmentFn< in: options.in, out: options.out, returnType: createIoSchema(options.out), - isEntry: true, + entryPoint: 'fragment', }; const call = ( @@ -170,9 +153,21 @@ export function fragmentFn< ...values: unknown[] ) => createFragmentFn(shell, stripTemplate(arg, ...values)); - return Object.assign(Object.assign(call, shell), { - does: call, - }) as TgpuFragmentFnShell; + return Object.assign(call, shell) as TgpuFragmentFnShell< + FragmentIn, + FragmentOut + >; +} + +export function isTgpuFragmentFn< + FragmentIn extends FragmentInConstrained, + FragmentOut extends FragmentOutConstrained, +>( + value: unknown | TgpuFragmentFn, +): value is TgpuFragmentFn { + return (value as TgpuFragmentFn)?.shell + ?.entryPoint === + 'fragment'; } // -------------- diff --git a/packages/typegpu/src/core/function/tgpuVertexFn.ts b/packages/typegpu/src/core/function/tgpuVertexFn.ts index d7a5c5e200..94a171a015 100644 --- a/packages/typegpu/src/core/function/tgpuVertexFn.ts +++ b/packages/typegpu/src/core/function/tgpuVertexFn.ts @@ -52,7 +52,7 @@ type TgpuVertexFnShellHeader< readonly in: VertexIn | undefined; readonly out: VertexOut; readonly argTypes: [IOLayoutToSchema] | []; - readonly isEntry: true; + readonly entryPoint: 'vertex'; }; /** @@ -77,19 +77,7 @@ export type TgpuVertexFnShell< & (( strings: TemplateStringsArray, ...values: unknown[] - ) => TgpuVertexFn, OmitBuiltins>) - & { - /** - * @deprecated Invoke the shell as a function instead. - */ - does: - & (( - implementation: (input: InferIO) => InferIO, - ) => TgpuVertexFn, OmitBuiltins>) - & (( - implementation: string, - ) => TgpuVertexFn, OmitBuiltins>); - }; + ) => TgpuVertexFn, OmitBuiltins>); export interface TgpuVertexFn< VertexIn extends VertexInConstrained = VertexInConstrained, @@ -146,7 +134,7 @@ export function vertexFn< argTypes: options.in && Object.keys(options.in).length !== 0 ? [createIoSchema(options.in)] : [], - isEntry: true, + entryPoint: 'vertex', }; const call = ( @@ -154,9 +142,17 @@ export function vertexFn< ...values: unknown[] ) => createVertexFn(shell, stripTemplate(arg, ...values)); - return Object.assign(Object.assign(call, shell), { - does: call, - }) as TgpuVertexFnShell; + return Object.assign(call, shell) as TgpuVertexFnShell; +} + +export function isTgpuVertexFn< + VertexIn extends VertexInConstrained, + VertexOut extends VertexOutConstrained, +>( + value: unknown | TgpuVertexFn, +): value is TgpuVertexFn { + return (value as TgpuVertexFn)?.shell?.entryPoint === + 'vertex'; } // -------------- diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index f803cdad65..03835df215 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -65,12 +65,16 @@ import type { ItemLayer, ItemStateStack, ResolutionCtx, + ShaderStage, TgpuShaderStage, Wgsl, } from './types.ts'; import { CodegenState, isSelfResolvable, NormalState } from './types.ts'; import type { WgslExtension } from './wgslExtensions.ts'; import { hasTinyestMetadata } from './shared/meta.ts'; +import { isTgpuComputeFn } from './core/function/tgpuComputeFn.ts'; +import { isTgpuVertexFn } from './core/function/tgpuVertexFn.ts'; +import { isTgpuFragmentFn } from './core/function/tgpuFragmentFn.ts'; /** * Inserted into bind group entry definitions that belong @@ -334,6 +338,12 @@ export class ResolutionCtxImpl implements ResolutionCtx { readonly #namespace: NamespaceInternal; readonly #shaderGenerator: ShaderGenerator; + /** + * Holds info about the currently resolved shader stage (if there is any). + * Note that if a function is used both in vertex and fragment stage, + * then it will only go through the process during the vertex stage. + */ + private _currentStage: ShaderStage; private readonly _indentController = new IndentController(); private readonly _itemStateStack = new ItemStateStackImpl(); readonly #modeStack: ExecState[] = []; @@ -444,6 +454,11 @@ export class ResolutionCtxImpl implements ResolutionCtx { } generateLog(op: string, args: Snippet[]): Snippet { + if (this._currentStage === 'vertex') { + throw new Error( + `'console.log' is not supported during vertex shader stage.`, + ); + } return this.#logGenerator.generateLog(this, op, args); } @@ -732,11 +747,24 @@ export class ResolutionCtxImpl implements ResolutionCtx { } if (isMarkedInternal(item) || hasTinyestMetadata(item)) { + // if we're resolving an entrypoint function, we want to update this._currentStage + const stage = getItemStage(item); + const resolutionAction = () => { + if (stage) { + this._currentStage = stage; + } + const result = this._getOrInstantiate(item); + if (stage) { + this._currentStage = undefined; + } + return result; + }; + // Top-level resolve if (this._itemStateStack.itemDepth === 0) { try { this.pushMode(new CodegenState()); - const result = provideCtx(this, () => this._getOrInstantiate(item)); + const result = provideCtx(this, resolutionAction); return snip( `${[...this._declarations].join('\n\n')}${result.value}`, Void, @@ -747,7 +775,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { } } - return this._getOrInstantiate(item); + return resolutionAction(); } // This is a value that comes from the outside, maybe we can coerce it @@ -972,3 +1000,15 @@ export function resolveFunctionHeader( } ` : `(${argList}) `; } + +function getItemStage(item: unknown): ShaderStage { + if (isTgpuComputeFn(item)) { + return 'compute'; + } + if (isTgpuVertexFn(item)) { + return 'vertex'; + } + if (isTgpuFragmentFn(item)) { + return 'fragment'; + } +} diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index 6d10ffc1d1..a6cf8f4213 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -217,6 +217,12 @@ export type ExecState = | CodegenState | SimulationState; +export type ShaderStage = + | 'vertex' + | 'fragment' + | 'compute' + | undefined; + /** * Passed into each resolvable item. All items in a tree share a resolution ctx, * but there can be layers added and removed from the item stack when going down diff --git a/packages/typegpu/tests/tgsl/consoleLog.test.ts b/packages/typegpu/tests/tgsl/consoleLog.test.ts index 1ebf4494c1..fc418aca89 100644 --- a/packages/typegpu/tests/tgsl/consoleLog.test.ts +++ b/packages/typegpu/tests/tgsl/consoleLog.test.ts @@ -10,13 +10,11 @@ import { it } from '../utils/extendedIt.ts'; describe('wgslGenerator with console.log', () => { let ctx: ResolutionCtxImpl; beforeEach(() => { - ctx = new ResolutionCtxImpl({ - namespace: namespace({ names: 'strict' }), - }); + ctx = new ResolutionCtxImpl({ namespace: namespace() }); ctx.pushMode(new CodegenState()); }); - it('Parses console.log in a stray function to a comment and warns', ({ root }) => { + it('Parses console.log in a stray function to a comment and warns', () => { using consoleWarnSpy = vi .spyOn(console, 'warn') .mockImplementation(() => {}); @@ -37,6 +35,100 @@ describe('wgslGenerator with console.log', () => { expect(consoleWarnSpy).toHaveBeenCalledTimes(1); }); + it('Does not reset the state after visiting one function', () => { + const fn1 = (n: number) => { + 'use gpu'; + }; + + const fn2 = (n: number) => { + 'use gpu'; + console.log(n); + }; + + const vs = tgpu['~unstable'].vertexFn({ out: { pos: d.builtin.position } })( + () => { + fn1(1); + fn2(2); + return { pos: d.vec4f() }; + }, + ); + expect(() => tgpu.resolve([vs])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - vertexFn:vs + - fn*:fn2(i32) + - fn:consoleLog: 'console.log' is not supported during vertex shader stage.] + `); + }); + + it('Throws appropriate error when in a vertex shader', () => { + const myLog = (n: number) => { + 'use gpu'; + console.log(n); + }; + + const vs = tgpu['~unstable'].vertexFn({ out: { pos: d.builtin.position } })( + () => { + myLog(5); + return { pos: d.vec4f() }; + }, + ); + expect(() => tgpu.resolve([vs])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - vertexFn:vs + - fn*:myLog(i32) + - fn:consoleLog: 'console.log' is not supported during vertex shader stage.] + `); + }); + + it('Throws appropriate error when in a vertex shader during pipeline resolution', ({ root }) => { + const myLog = (n: number) => { + 'use gpu'; + console.log(n); + }; + + const vs = tgpu['~unstable'].vertexFn({ out: { pos: d.builtin.position } })( + () => { + myLog(5); + return { pos: d.vec4f() }; + }, + ); + const fs = tgpu['~unstable'].fragmentFn({ out: d.vec4f })(() => { + return d.vec4f(); + }); + + const pipeline = root['~unstable'] + .withVertex(vs) + .withFragment(fs, { format: 'rg8unorm' }) + .createPipeline(); + + expect(() => tgpu.resolve([pipeline])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - renderPipeline:pipeline + - renderPipelineCore + - vertexFn:vs + - fn*:myLog(i32) + - fn:consoleLog: 'console.log' is not supported during vertex shader stage.] + `); + }); + + it('Ignores console.log in a fragment shader resolved without a pipeline', () => { + const fs = tgpu['~unstable'] + .fragmentFn({ out: d.vec4f })(() => { + console.log(d.u32(321)); + return d.vec4f(); + }); + + expect(tgpu.resolve([fs])).toMatchInlineSnapshot(` + "@fragment fn fs() -> @location(0) vec4f { + /* console.log() */; + return vec4f(); + }" + `); + }); + it('Parses a single console.log in a render pipeline', ({ root }) => { const vs = tgpu['~unstable'] .vertexFn({ out: { pos: d.builtin.position } })(() => { @@ -49,7 +141,7 @@ describe('wgslGenerator with console.log', () => { }); const pipeline = root['~unstable'] - .withVertex(vs, {}) + .withVertex(vs) .withFragment(fs, { format: 'rg8unorm' }) .createPipeline();