feat: Make tgpu.declare dereferenceable

apps/typegpu-docs/src/content/docs/apis/utils.mdx

@@ -219,14 +219,22 @@ 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);',
 );
 
 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:

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<string, unknown>): this;
+  readonly $: InferGPU<Void>;
+  readonly [$gpuValueOf]: InferGPU<Void>;
 }
 
 /**
@@ -31,15 +36,54 @@ 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;
 
   constructor(declaration: string) {
     this.#declaration = declaration;
   }
 
+  // prototype properties
+  declare [$internal]: true;
+  declare $: InferGPU<Void>;
+  declare readonly [$gpuValueOf]: InferGPU<Void>;
+
+  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<string, unknown>): 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}`;
-  }
 }

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<string, never>,
+  in Varying extends TgpuFragmentFn.In = never,
   out Output extends TgpuFragmentFn.Out = TgpuFragmentFn.Out,
 > extends TgpuNamable {
   readonly [$internal]: true;

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<string, never>,
+  in VertexIn extends TgpuVertexFn.In = never,
   out VertexOut extends TgpuVertexFn.Out = TgpuVertexFn.Out,
 > extends TgpuNamable {
   readonly [$internal]: true;

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<storage, read_write> 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<storage, read> 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<uniform> buf: i32;
+      fn helper() -> i32 {
+        return buf;
+      }"
+    `);
+  });
 });

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<TgpuComputeFn>();
+  });
+
+  it('compute with in is assignable to a bare TgpuComputeFn', () => {
+    const computeMain = tgpu.computeFn({
+      in: { gid: d.builtin.globalInvocationId },
+      workgroupSize: [1],
+    })``;
+
+    expectTypeOf(computeMain).toExtend<TgpuComputeFn>();
+  });
+});

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<uniform> 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<uniform> val: f32;');
+
+    // oxlint-disable-next-line typescript/no-base-to-string
+    expect(`${declaration}`).toMatchInlineSnapshot(
+      `"declare: @group(0) @binding(0) var<uniform> val: f32;"`,
+    );
+  });
+
   it("throws when '$uses' is called multiple times", () => {
     const declaration = tgpu['~unstable']
       .declare('@group(0) @binding(0) var<myStruct> 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;
+      }"
+    `);
+  });
 });

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<TgpuFragmentFn>();
+  });
+
+  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<TgpuFragmentFn>();
+  });
+
+  it('vertex is assignable to a bare TgpuVertexFn', () => {
+    const vertexMain = tgpu.vertexFn({
+      out: { pos: d.builtin.position },
+    })``;
+
+    expectTypeOf(vertexMain).toExtend<TgpuVertexFn>();
+  });
+
+  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<TgpuVertexFn>();
+  });
+});