diff --git a/apps/typegpu-docs/src/content/docs/apis/utils.mdx b/apps/typegpu-docs/src/content/docs/apis/utils.mdx index 223406342d..89225ebb07 100644 --- a/apps/typegpu-docs/src/content/docs/apis/utils.mdx +++ b/apps/typegpu-docs/src/content/docs/apis/utils.mdx @@ -219,7 +219,9 @@ allowing functions to branch based on what language is being generated. `tgpu['~unstable'].declare(source)` adds a declaration to generated WGSL whenever a resolving object depends on it. It is an escape hatch for WGSL declarations that TypeGPU cannot yet express directly: -```ts +```ts twoslash +import { tgpu } from 'typegpu'; +// ---cut--- const diagnosticDirective = tgpu['~unstable'].declare( 'diagnostic(off, derivative_uniformity);', ); @@ -227,6 +229,12 @@ const diagnosticDirective = tgpu['~unstable'].declare( const helper = tgpu.fn([])`() { // ... }`.$uses({ diagnosticDirective }); + +const anotherHelper = () => { + 'use gpu'; + // The declaration will not affect the function's body. + diagnosticDirective.$; +} ``` Declarations can reference TypeGPU resources through one dependency map: diff --git a/packages/typegpu/src/core/declare/tgpuDeclare.ts b/packages/typegpu/src/core/declare/tgpuDeclare.ts index a105518138..a309314479 100644 --- a/packages/typegpu/src/core/declare/tgpuDeclare.ts +++ b/packages/typegpu/src/core/declare/tgpuDeclare.ts @@ -1,8 +1,10 @@ -import { type ResolvedSnippet, snip } from '../../data/snippet.ts'; +import { snip } from '../../data/snippet.ts'; import { Void } from '../../data/wgslTypes.ts'; -import { $internal, $resolve } from '../../shared/symbols.ts'; -import type { ResolutionCtx, SelfResolvable } from '../../types.ts'; +import { $gpuValueOf, $internal } from '../../shared/symbols.ts'; import { type ExternalMap, replaceExternalsInWgsl } from '../resolve/externals.ts'; +import type { InferGPU } from '../../data/index.ts'; +import { makeDereferenceable } from '../../tgsl/makeDereferenceable.ts'; +import { makeResolvable } from '../../tgsl/makeResolvable.ts'; // ---------- // Public API @@ -13,7 +15,10 @@ import { type ExternalMap, replaceExternalsInWgsl } from '../resolve/externals.t * when resolving objects that use it. */ export interface TgpuDeclare { + readonly [$internal]: true; $uses(dependencyMap: Record): this; + readonly $: InferGPU; + readonly [$gpuValueOf]: InferGPU; } /** @@ -31,8 +36,7 @@ export function declare(declaration: string): TgpuDeclare { // Implementation // -------------- -export class TgpuDeclareImpl implements TgpuDeclare, SelfResolvable { - readonly [$internal] = true; +export class TgpuDeclareImpl implements TgpuDeclare { #externals: ExternalMap | undefined; #declaration: string; @@ -40,6 +44,46 @@ export class TgpuDeclareImpl implements TgpuDeclare, SelfResolvable { this.#declaration = declaration; } + // prototype properties + declare [$internal]: true; + declare $: InferGPU; + declare readonly [$gpuValueOf]: InferGPU; + + static { + TgpuDeclareImpl.prototype[$internal] = true; + + makeDereferenceable( + makeResolvable(TgpuDeclareImpl.prototype, { + asString() { + return `declare: ${this.#declaration}`; + }, + resolve(ctx) { + const replacedDeclaration = replaceExternalsInWgsl( + ctx, + this.#externals ?? {}, + this.#declaration, + ); + + ctx.addDeclaration(replacedDeclaration); + return snip('', Void, /* origin */ 'constant'); + }, + }), + { + codegenMode: { + getBaseSnippet(trackingProxy) { + return snip( + trackingProxy, + Void, + /* origin */ 'constant', + /* possibleSideEffects */ false, + ); + }, + }, + normalMode: { get() {} }, + }, + ); + } + $uses(dependencyMap: Record): this { if (this.#externals !== undefined) { throw new Error( @@ -49,19 +93,4 @@ export class TgpuDeclareImpl implements TgpuDeclare, SelfResolvable { this.#externals = dependencyMap; return this; } - - [$resolve](ctx: ResolutionCtx): ResolvedSnippet { - const replacedDeclaration = replaceExternalsInWgsl( - ctx, - this.#externals ?? {}, - this.#declaration, - ); - - ctx.addDeclaration(replacedDeclaration); - return snip('', Void, /* origin */ 'constant'); - } - - toString() { - return `declare: ${this.#declaration}`; - } } diff --git a/packages/typegpu/src/core/function/tgpuFragmentFn.ts b/packages/typegpu/src/core/function/tgpuFragmentFn.ts index e6d3b8091f..a909e4d733 100644 --- a/packages/typegpu/src/core/function/tgpuFragmentFn.ts +++ b/packages/typegpu/src/core/function/tgpuFragmentFn.ts @@ -107,7 +107,7 @@ export interface TgpuFragmentFnShell< export interface TgpuFragmentFn< // @ts-expect-error: We override the variance - in Varying extends TgpuFragmentFn.In = Record, + in Varying extends TgpuFragmentFn.In = never, out Output extends TgpuFragmentFn.Out = TgpuFragmentFn.Out, > extends TgpuNamable { readonly [$internal]: true; diff --git a/packages/typegpu/src/core/function/tgpuVertexFn.ts b/packages/typegpu/src/core/function/tgpuVertexFn.ts index e0d2d4c388..877a053bbc 100644 --- a/packages/typegpu/src/core/function/tgpuVertexFn.ts +++ b/packages/typegpu/src/core/function/tgpuVertexFn.ts @@ -78,7 +78,7 @@ export interface TgpuVertexFnShell< export interface TgpuVertexFn< // @ts-expect-error: We override the variance - in VertexIn extends TgpuVertexFn.In = Record, + in VertexIn extends TgpuVertexFn.In = never, out VertexOut extends TgpuVertexFn.Out = TgpuVertexFn.Out, > extends TgpuNamable { readonly [$internal]: true; diff --git a/packages/typegpu/tests/bufferBindings.test.ts b/packages/typegpu/tests/bufferBindings.test.ts index 770bb81951..ee9af57660 100644 --- a/packages/typegpu/tests/bufferBindings.test.ts +++ b/packages/typegpu/tests/bufferBindings.test.ts @@ -230,6 +230,25 @@ describe('root.createMutable', () => { mappedBuffer.write([{ id: 1, values: d.vec3f() }]); }); }); + + it('should not require .$ in template resolve', ({ root }) => { + const buf = root.createMutable(d.i32); + + expect( + tgpu.resolve({ + externals: { buf }, + template: ` +fn helper() -> i32 { + return buf; +}`, + }), + ).toMatchInlineSnapshot(` + "@group(0) @binding(0) var buf: i32; + fn helper() -> i32 { + return buf; + }" + `); + }); }); describe('root.createReadonly', () => { @@ -498,6 +517,25 @@ describe('root.createReadonly', () => { mappedBuffer.write([{ id: 1, values: d.vec3f() }]); }); }); + + it('should not require .$ in template resolve', ({ root }) => { + const buf = root.createReadonly(d.i32); + + expect( + tgpu.resolve({ + externals: { buf }, + template: ` +fn helper() -> i32 { + return buf; +}`, + }), + ).toMatchInlineSnapshot(` + "@group(0) @binding(0) var buf: i32; + fn helper() -> i32 { + return buf; + }" + `); + }); }); describe('root.createUniform', () => { @@ -728,4 +766,23 @@ describe('root.createUniform', () => { mappedBuffer.write([{ id: 1, values: d.vec3f() }]); }); }); + + it('should not require .$ in template resolve', ({ root }) => { + const buf = root.createUniform(d.i32); + + expect( + tgpu.resolve({ + externals: { buf }, + template: ` +fn helper() -> i32 { + return buf; +}`, + }), + ).toMatchInlineSnapshot(` + "@group(0) @binding(0) var buf: i32; + fn helper() -> i32 { + return buf; + }" + `); + }); }); diff --git a/packages/typegpu/tests/computePipeline.test.ts b/packages/typegpu/tests/computePipeline.test.ts index 3d68551d45..fb9326cf7f 100644 --- a/packages/typegpu/tests/computePipeline.test.ts +++ b/packages/typegpu/tests/computePipeline.test.ts @@ -1,5 +1,12 @@ import { describe, expect, expectTypeOf, vi } from 'vitest'; -import { d, isBindGroup, MissingBindGroupsError, tgpu, type TgpuComputePipeline } from 'typegpu'; +import { + d, + isBindGroup, + MissingBindGroupsError, + tgpu, + type TgpuComputeFn, + type TgpuComputePipeline, +} from 'typegpu'; import { restoreResource, snapshotResource } from 'typegpu/~internal'; import { it } from 'typegpu-testing-utility'; import { extensionEnabled } from 'typegpu/std'; @@ -840,3 +847,22 @@ describe('pipe', () => { expect(pipeline.pipe((p) => p)).toBe(pipeline); }); }); + +describe('types', () => { + it('compute is assignable to a bare TgpuComputeFn', () => { + const computeMain = tgpu.computeFn({ + workgroupSize: [1], + })``; + + expectTypeOf(computeMain).toExtend(); + }); + + it('compute with in is assignable to a bare TgpuComputeFn', () => { + const computeMain = tgpu.computeFn({ + in: { gid: d.builtin.globalInvocationId }, + workgroupSize: [1], + })``; + + expectTypeOf(computeMain).toExtend(); + }); +}); diff --git a/packages/typegpu/tests/declare.test.ts b/packages/typegpu/tests/declare.test.ts index cd5d95679f..64cb00b138 100644 --- a/packages/typegpu/tests/declare.test.ts +++ b/packages/typegpu/tests/declare.test.ts @@ -93,7 +93,7 @@ struct Output { [], d.f32, )(() => { - declaration; + declaration.$; return 2; }); @@ -106,6 +106,27 @@ struct Output { `); }); + it('is a noop in JS', () => { + const declaration = tgpu['~unstable'].declare('@group(0) @binding(0) var val: f32;'); + + const main = () => { + 'use gpu'; + declaration.$; + return 2 + 2; + }; + + expect(main()).toBe(4); + }); + + it('has a readable toString', () => { + const declaration = tgpu['~unstable'].declare('@group(0) @binding(0) var val: f32;'); + + // oxlint-disable-next-line typescript/no-base-to-string + expect(`${declaration}`).toMatchInlineSnapshot( + `"declare: @group(0) @binding(0) var val: f32;"`, + ); + }); + it("throws when '$uses' is called multiple times", () => { const declaration = tgpu['~unstable'] .declare('@group(0) @binding(0) var val: f32;') @@ -117,4 +138,25 @@ struct Output { `[Error: Cannot call '$uses' multiple times. If you wish to override dependencies, use slots or accessors instead.]`, ); }); + + it('should not require .$ in template resolve', () => { + const declaration = tgpu['~unstable'].declare('const C = 4;'); + + expect( + tgpu.resolve({ + externals: { declaration }, + template: ` +declaration +fn helper() -> i32 { + return C; +}`, + }), + ).toMatchInlineSnapshot(` + "const C = 4; + + fn helper() -> i32 { + return C; + }" + `); + }); }); diff --git a/packages/typegpu/tests/renderPipeline.test.ts b/packages/typegpu/tests/renderPipeline.test.ts index b9876c0009..5f454b374f 100644 --- a/packages/typegpu/tests/renderPipeline.test.ts +++ b/packages/typegpu/tests/renderPipeline.test.ts @@ -2460,3 +2460,39 @@ describe('pipe', () => { expectTypeOf(pipeline.pipe((p) => p)).toEqualTypeOf(pipeline); }); }); + +describe('types', () => { + it('fragment is assignable to a bare TgpuFragmentFn', () => { + const fragmentMain = tgpu.fragmentFn({ + out: d.vec4f, + })``; + + expectTypeOf(fragmentMain).toExtend(); + }); + + it('fragment with in is assignable to a bare TgpuFragmentFn', () => { + const fragmentMain = tgpu.fragmentFn({ + in: { coordinates: d.location(4, d.vec2i) }, + out: d.vec4f, + })``; + + expectTypeOf(fragmentMain).toExtend(); + }); + + it('vertex is assignable to a bare TgpuVertexFn', () => { + const vertexMain = tgpu.vertexFn({ + out: { pos: d.builtin.position }, + })``; + + expectTypeOf(vertexMain).toExtend(); + }); + + it('vertex with in is assignable to a bare TgpuVertexFn', () => { + const vertexMain = tgpu.vertexFn({ + in: { coordinates: d.location(4, d.vec2i) }, + out: { pos: d.builtin.position }, + })``; + + expectTypeOf(vertexMain).toExtend(); + }); +});