fix: Remove redundant .with slot calls & chore: Cleanup ItemStateStack (#1853)

Author: aleksanderkatan

apps/typegpu-docs/src/examples/simulation/fluid-double-buffering/index.ts

@@ -454,7 +454,6 @@ function makePipelines(
   outputGridMutable: TgpuBufferMutable<GridData>,
 ) {
   const initWorldPipeline = root['~unstable']
-    .with(inputGridSlot, outputGridMutable)
     .with(outputGridSlot, outputGridMutable)
     .createGuardedComputePipeline((xu, yu) => {
       'use gpu';

packages/typegpu/src/core/function/tgpuFn.ts

@@ -291,7 +291,7 @@ function createBoundFunction<ImplSchema extends AnyFn>(
       slot: TgpuSlot<unknown> | TgpuAccessor,
       value: unknown,
     ): TgpuFn<ImplSchema> {
-      return createBoundFunction(fn, [
+      return createBoundFunction(innerFn, [
         ...pairs,
         [isAccessor(slot) ? slot.slot : slot, value],
       ]);

packages/typegpu/src/resolutionCtx.ts

@@ -76,6 +76,7 @@ import type {
   ItemLayer,
   ItemStateStack,
   ResolutionCtx,
+  StackLayer,
   TgpuShaderStage,
   Wgsl,
 } from './types.ts';
@@ -103,23 +104,8 @@ export type ResolutionCtxImplOptions = {
   readonly namespace: Namespace;
 };
 
-type SlotBindingLayer = {
-  type: 'slotBinding';
-  bindingMap: WeakMap<TgpuSlot<unknown>, unknown>;
-};
-
-type BlockScopeLayer = {
-  type: 'blockScope';
-  declarations: Map<string, Snippet>;
-};
-
 class ItemStateStackImpl implements ItemStateStack {
-  private _stack: (
-    | ItemLayer
-    | SlotBindingLayer
-    | FunctionScopeLayer
-    | BlockScopeLayer
-  )[] = [];
+  private _stack: StackLayer[] = [];
   private _itemDepth = 0;
 
   get itemDepth(): number {
@@ -146,21 +132,13 @@ class ItemStateStackImpl implements ItemStateStack {
     });
   }
 
-  popItem() {
-    this.pop('item');
-  }
-
   pushSlotBindings(pairs: SlotValuePair<unknown>[]) {
     this._stack.push({
       type: 'slotBinding',
       bindingMap: new WeakMap(pairs),
     });
   }
 
-  popSlotBindings() {
-    this.pop('slotBinding');
-  }
-
   pushFunctionScope(
     functionType: 'normal' | TgpuShaderStage,
     args: Snippet[],
@@ -182,31 +160,26 @@ class ItemStateStackImpl implements ItemStateStack {
     return scope;
   }
 
-  popFunctionScope() {
-    this.pop('functionScope');
-  }
-
   pushBlockScope() {
     this._stack.push({
       type: 'blockScope',
       declarations: new Map(),
     });
   }
 
-  popBlockScope() {
-    this.pop('blockScope');
-  }
-
-  pop(type?: (typeof this._stack)[number]['type']) {
+  pop<T extends StackLayer['type']>(type: T): Extract<StackLayer, { type: T }>;
+  pop(): StackLayer | undefined;
+  pop(type?: StackLayer['type']) {
     const layer = this._stack[this._stack.length - 1];
     if (!layer || (type && layer.type !== type)) {
       throw new Error(`Internal error, expected a ${type} layer to be on top.`);
     }
 
-    this._stack.pop();
+    const poppedValue = this._stack.pop();
     if (type === 'item') {
       this._itemDepth--;
     }
+    return poppedValue;
   }
 
   readSlot<T>(slot: TgpuSlot<T>): T | undefined {
@@ -452,7 +425,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
   }
 
   popBlockScope() {
-    this._itemStateStack.popBlockScope();
+    this._itemStateStack.pop('blockScope');
   }
 
   generateLog(op: string, args: Snippet[]): Snippet {
@@ -560,7 +533,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
       };
     } finally {
       if (fnScopePushed) {
-        this._itemStateStack.popFunctionScope();
+        this._itemStateStack.pop('functionScope');
       }
       this.#namespaceInternal.nameRegistry.popFunctionScope();
     }
@@ -612,7 +585,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     try {
       return callback();
     } finally {
-      this._itemStateStack.popSlotBindings();
+      this._itemStateStack.pop('slotBinding');
     }
   }
 
@@ -700,7 +673,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
 
       throw new ResolutionError(err, [derived]);
     } finally {
-      this._itemStateStack.popItem();
+      this._itemStateStack.pop('item');
     }
   }
 
@@ -772,7 +745,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
 
       throw new ResolutionError(err, [item]);
     } finally {
-      this._itemStateStack.popItem();
+      this._itemStateStack.pop('item');
     }
   }
 

packages/typegpu/src/types.ts

@@ -109,15 +109,29 @@ export type FunctionScopeLayer = {
   reportedReturnTypes: Set<AnyData>;
 };
 
+export type SlotBindingLayer = {
+  type: 'slotBinding';
+  bindingMap: WeakMap<TgpuSlot<unknown>, unknown>;
+};
+
+export type BlockScopeLayer = {
+  type: 'blockScope';
+  declarations: Map<string, Snippet>;
+};
+
+export type StackLayer =
+  | ItemLayer
+  | SlotBindingLayer
+  | FunctionScopeLayer
+  | BlockScopeLayer;
+
 export interface ItemStateStack {
   readonly itemDepth: number;
   readonly topItem: ItemLayer;
   readonly topFunctionScope: FunctionScopeLayer | undefined;
 
   pushItem(): void;
-  popItem(): void;
   pushSlotBindings(pairs: SlotValuePair<unknown>[]): void;
-  popSlotBindings(): void;
   pushFunctionScope(
     functionType: 'normal' | TgpuShaderStage,
     args: Snippet[],
@@ -129,10 +143,11 @@ export interface ItemStateStack {
     returnType: AnyData | undefined,
     externalMap: Record<string, unknown>,
   ): FunctionScopeLayer;
-  popFunctionScope(): void;
   pushBlockScope(): void;
-  popBlockScope(): void;
-  pop(type?: 'functionScope' | 'blockScope' | 'slotBinding' | 'item'): void;
+
+  pop<T extends StackLayer['type']>(type: T): Extract<StackLayer, { type: T }>;
+  pop(): StackLayer | undefined;
+
   readSlot<T>(slot: TgpuSlot<T>): T | undefined;
   getSnippetById(id: string): Snippet | undefined;
   defineBlockVariable(id: string, snippet: Snippet): void;

packages/typegpu/tests/slot.test.ts

@@ -1,7 +1,7 @@
 import { describe, expect } from 'vitest';
 import * as d from '../src/data/index.ts';
-import * as std from '../src/std/index.ts';
 import tgpu from '../src/index.ts';
+import * as std from '../src/std/index.ts';
 import { it } from './utils/extendedIt.ts';
 
 const RED = 'vec3f(1., 0., 0.)';
@@ -363,4 +363,32 @@ describe('tgpu.slot', () => {
       }"
     `);
   });
+
+  it('includes slot bindings in toString', () => {
+    const firstSlot = tgpu.slot<number>();
+    const secondSlot = tgpu.slot<number>();
+    const thirdSlot = tgpu.slot<number>();
+
+    const getSize = tgpu.fn([], d.f32)(() =>
+      firstSlot.$ + secondSlot.$ + thirdSlot.$
+    )
+      .with(firstSlot, 1)
+      .with(secondSlot, 2)
+      .with(thirdSlot, 3);
+
+    expect(getSize.toString()).toMatchInlineSnapshot(
+      `"fn:getSize[firstSlot=1, secondSlot=2, thirdSlot=3]"`,
+    );
+  });
+
+  it('safe stringifies in toString', () => {
+    const slot = tgpu.slot<d.v4f>();
+
+    const getSize = tgpu.fn([], d.f32)(() => slot.$.x)
+      .with(slot, d.vec4f(1, 2, 3, 4));
+
+    expect(getSize.toString()).toMatchInlineSnapshot(
+      `"fn:getSize[slot=vec4f(1, 2, 3, 4)]"`,
+    );
+  });
 });

packages/typegpu/tests/tgsl/wgslGenerator.test.ts

@@ -338,7 +338,7 @@ describe('wgslGenerator', () => {
       const res2 = wgslGenerator.expression(
         (astInfo.ast?.body[1][1] as tinyest.Const)[2],
       );
-      ctx[$internal].itemStateStack.popBlockScope();
+      ctx[$internal].itemStateStack.pop('blockScope');
 
       expect(res2.dataType).toStrictEqual(d.vec4f);
 
@@ -353,7 +353,7 @@ describe('wgslGenerator', () => {
       const res4 = wgslGenerator.expression(
         astInfo.ast?.body[1][2] as tinyest.Expression,
       );
-      ctx[$internal].itemStateStack.popBlockScope();
+      ctx[$internal].itemStateStack.pop('blockScope');
 
       expect(res3.dataType).toStrictEqual(d.atomic(d.u32));
       expect(res4.dataType).toStrictEqual(Void);

packages/typegpu/tests/utils/parseResolved.ts

@@ -58,8 +58,8 @@ function extractSnippetFromFn(cb: () => unknown): Snippet {
       } finally {
         if (pushedFnScope) {
           ctx.popBlockScope();
-          ctx[$internal].itemStateStack.popFunctionScope();
-          ctx[$internal].itemStateStack.popItem();
+          ctx[$internal].itemStateStack.pop('functionScope');
+          ctx[$internal].itemStateStack.pop('item');
         }
         ctx.popMode('codegen');
       }