feat: Expected type stack in WgslGenerator (#1532)

* Add struct default constructor

* Update docs

* Update packages/typegpu/tests/struct.test.ts

Co-authored-by: Iwo Plaza <[email protected]>

* Update packages/typegpu/src/tgsl/wgslGenerator.ts

Co-authored-by: Konrad Reczko <[email protected]>

* Change [1] to [0]

* Make array schemas callable (but they do nothing)

* Move the struct resolve test to wgslGenerator.test.ts

* Move tests back to struct.test.ts

* Add array tests

* Implement JS side of array calls

* Implement WGSL side

* Change u32 to i32

* Add tgsl parsing tests

* Remove unnecessary (I hope so) parentheses

* Docs

* Lint

* Fix import

* Lint

* Nits

* Remove generic from struct constructor, update docs

* Apply suggestions from code review

Co-authored-by: Iwo Plaza <[email protected]>

* Update array toString, override dualImpl toString

* Add some tests

* Add expectedTypeStack

* Use expectedTypeStack in struct call

* Use expectedTypeStack for returns and nested structs

* Remove callStack from return value generation, comment out some weird tests

* Remove the arbitrary object cast to output type

* CallStack is kil

* Remove legacy calls to stack from tgpuFn

* Remove legacy calls from fragmentFn and tgpuFn

* Add array support

* Add type coercion, apply type coercion to return value

* Merge fixes

* Move struct/array handling

* Cleanup struct/array call

* Refactor

* Add function support

* Add bool to the stack for if/for/while conditions

* Bring back commented test 1

* Bring back commented tests 2 & 3

* Refactor array generation, add better comments

* Rename `argTypes` to `argConversionHint`

* Remove UnknownData from stack type

* Docs for expectedTypeStack

* Lint

* Update error messages

* Nits

* Better convert to null message

* Update tests to still create Output structs before return statement

* Unify type errors

* Simpler tests, review fixes

* Holding the expected type stack... on the stack

* Tweaks

* Review fixes

* Shuffle things around

---------

Co-authored-by: Aleksander Katan <[email protected]>
Co-authored-by: Iwo Plaza <[email protected]>
Co-authored-by: Konrad Reczko <[email protected]>

Author: 4 people

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

@@ -0,0 +1,26 @@
+import type { TgpuDualFn } from '../../data/dataTypes.ts';
+import type { MapValueToSnippet, Snippet } from '../../data/snippet.ts';
+import { inCodegenMode } from '../../execMode.ts';
+import type { FnArgsConversionHint } from '../../types.ts';
+import { setName } from '../../shared/meta.ts';
+import { $internal } from '../../shared/symbols.ts';
+
+export function createDualImpl<T extends (...args: never[]) => unknown>(
+  jsImpl: T,
+  gpuImpl: (...args: MapValueToSnippet<Parameters<T>>) => Snippet,
+  name: string,
+  argConversionHint: FnArgsConversionHint = 'keep',
+): TgpuDualFn<T> {
+  const impl = ((...args: Parameters<T>) => {
+    if (inCodegenMode()) {
+      return gpuImpl(...(args as MapValueToSnippet<Parameters<T>>)) as Snippet;
+    }
+    return jsImpl(...args);
+  }) as T;
+
+  setName(impl, name);
+  impl.toString = () => name;
+  (impl as TgpuDualFn<T>)[$internal] = { jsImpl, gpuImpl, argConversionHint };
+
+  return impl as TgpuDualFn<T>;
+}

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

@@ -4,7 +4,6 @@ import { snip } from '../../data/snippet.ts';
 import { Void } from '../../data/wgslTypes.ts';
 import { ExecutionError } from '../../errors.ts';
 import { provideInsideTgpuFn } from '../../execMode.ts';
-import { createDualImpl } from '../../shared/generators.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
 import { getName, setName } from '../../shared/meta.ts';
 import type { Infer } from '../../shared/repr.ts';
@@ -14,7 +13,6 @@ import {
   $providing,
 } from '../../shared/symbols.ts';
 import type { Prettify } from '../../shared/utilityTypes.ts';
-import type { GenerationCtx } from '../../tgsl/generationHelpers.ts';
 import type {
   FnArgsConversionHint,
   ResolutionCtx,
@@ -43,6 +41,7 @@ import type {
   InheritArgNames,
 } from './fnTypes.ts';
 import { stripTemplate } from './templateUtils.ts';
+import { createDualImpl } from './dualImpl.ts';
 
 // ----------
 // Public API
@@ -85,7 +84,7 @@ export type TgpuFnShell<
 interface TgpuFnBase<ImplSchema extends AnyFn> extends TgpuNamable {
   readonly [$internal]: {
     implementation: Implementation<ImplSchema>;
-    argTypes: FnArgsConversionHint;
+    argConversionHint: FnArgsConversionHint;
   };
   readonly resourceType: 'function';
   readonly shell: TgpuFnShellHeader<
@@ -170,7 +169,7 @@ function createFn<ImplSchema extends AnyFn>(
   const fnBase: This = {
     [$internal]: {
       implementation,
-      argTypes: shell.argTypes,
+      argConversionHint: shell.argTypes,
     },
     shell,
     resourceType: 'function' as const,
@@ -207,23 +206,9 @@ function createFn<ImplSchema extends AnyFn>(
           shell.returnType,
           core.applyExternals,
         );
-
-        return core.resolve(ctx, shell.argTypes, shell.returnType);
       }
 
-      const generationCtx = ctx as GenerationCtx;
-      if (generationCtx.callStack === undefined) {
-        throw new Error(
-          'Cannot resolve a TGSL function outside of a generation context',
-        );
-      }
-
-      try {
-        generationCtx.callStack.push(shell.returnType);
-        return core.resolve(ctx, shell.argTypes, shell.returnType);
-      } finally {
-        generationCtx.callStack.pop();
-      }
+      return core.resolve(ctx, shell.argTypes, shell.returnType);
     },
   };
 
@@ -284,7 +269,7 @@ function createBoundFunction<ImplSchema extends AnyFn>(
   const fnBase: This = {
     [$internal]: {
       implementation: innerFn[$internal].implementation,
-      argTypes: innerFn[$internal].argTypes,
+      argConversionHint: innerFn[$internal].argConversionHint,
     },
     resourceType: 'function',
     shell: innerFn.shell,

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

@@ -17,7 +17,6 @@ import {
   type TgpuNamable,
 } from '../../shared/meta.ts';
 import { $getNameForward, $internal } from '../../shared/symbols.ts';
-import type { GenerationCtx } from '../../tgsl/generationHelpers.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { addReturnTypeToExternals } from '../resolve/externals.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
@@ -227,31 +226,11 @@ function createFragmentFn(
       }
       core.applyExternals({ Out: outputType });
 
-      if (typeof implementation === 'string') {
-        return core.resolve(
-          ctx,
-          inputWithLocation ? [inputWithLocation] : [],
-          shell.returnType,
-        );
-      }
-
-      const generationCtx = ctx as GenerationCtx;
-      if (generationCtx.callStack === undefined) {
-        throw new Error(
-          'Cannot resolve a TGSL function outside of a generation context',
-        );
-      }
-
-      try {
-        generationCtx.callStack.push(outputType);
-        return core.resolve(
-          ctx,
-          inputWithLocation ? [inputWithLocation] : [],
-          shell.returnType,
-        );
-      } finally {
-        generationCtx.callStack.pop();
-      }
+      return core.resolve(
+        ctx,
+        inputWithLocation ? [inputWithLocation] : [],
+        shell.returnType,
+      );
     },
 
     toString() {

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

@@ -16,7 +16,6 @@ import {
   type TgpuNamable,
 } from '../../shared/meta.ts';
 import { $getNameForward, $internal } from '../../shared/symbols.ts';
-import type { GenerationCtx } from '../../tgsl/generationHelpers.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
 import type {
@@ -207,31 +206,13 @@ function createVertexFn(
           core.applyExternals({ In: inputType });
         }
         core.applyExternals({ Out: outputWithLocation });
-
-        return core.resolve(
-          ctx,
-          shell.argTypes,
-          outputWithLocation,
-        );
-      }
-
-      const generationCtx = ctx as GenerationCtx;
-      if (generationCtx.callStack === undefined) {
-        throw new Error(
-          'Cannot resolve a TGSL function outside of a generation context',
-        );
       }
 
-      try {
-        generationCtx.callStack.push(outputWithLocation);
-        return core.resolve(
-          ctx,
-          shell.argTypes,
-          outputWithLocation,
-        );
-      } finally {
-        generationCtx.callStack.pop();
-      }
+      return core.resolve(
+        ctx,
+        shell.argTypes,
+        outputWithLocation,
+      );
     },
 
     toString() {

packages/typegpu/src/data/dataTypes.ts

@@ -28,7 +28,7 @@ export type TgpuDualFn<TImpl extends (...args: never[]) => unknown> =
     [$internal]: {
       jsImpl: TImpl | string;
       gpuImpl: (...args: MapValueToSnippet<Parameters<TImpl>>) => Snippet;
-      argTypes: FnArgsConversionHint;
+      argConversionHint: FnArgsConversionHint;
     };
   };
 

packages/typegpu/src/data/matrix.ts

@@ -1,4 +1,4 @@
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 import type { $repr } from '../shared/symbols.ts';
 import { $internal } from '../shared/symbols.ts';
 import type { SelfResolvable } from '../types.ts';

packages/typegpu/src/data/numeric.ts

@@ -1,4 +1,4 @@
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 import { $internal } from '../shared/symbols.ts';
 import { snip } from './snippet.ts';
 import type {

packages/typegpu/src/data/vector.ts

@@ -1,4 +1,4 @@
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 import { $repr } from '../shared/symbols.ts';
 import { snip } from './snippet.ts';
 import { bool, f16, f32, i32, u32 } from './numeric.ts';

packages/typegpu/src/errors.ts

@@ -197,3 +197,12 @@ export class IllegalBufferAccessError extends Error {
     Object.setPrototypeOf(this, IllegalBufferAccessError.prototype);
   }
 }
+
+export class WgslTypeError extends Error {
+  constructor(msg: string) {
+    super(msg);
+
+    // Set the prototype explicitly.
+    Object.setPrototypeOf(this, WgslTypeError.prototype);
+  }
+}

packages/typegpu/src/resolutionCtx.ts

@@ -101,6 +101,14 @@ class ItemStateStackImpl implements ItemStateStack {
     return state;
   }
 
+  get topFunctionReturnType(): AnyData {
+    const scope = this._stack.findLast((e) => e.type === 'functionScope');
+    if (!scope) {
+      throw new Error('Internal error, expected function scope to be present.');
+    }
+    return scope.returnType;
+  }
+
   pushItem() {
     this._itemDepth++;
     this._stack.push({
@@ -343,8 +351,8 @@ export class ResolutionCtxImpl implements ResolutionCtx {
   public readonly fixedBindings: FixedBindingConfig[] = [];
   // --
 
-  public readonly callStack: unknown[] = [];
   public readonly names: NameRegistry;
+  public expectedType: AnyData | undefined;
 
   constructor(opts: ResolutionCtxImplOptions) {
     this.names = opts.names;
@@ -354,6 +362,10 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     return this._indentController.pre;
   }
 
+  get topFunctionReturnType() {
+    return this._itemStateStack.topFunctionReturnType;
+  }
+
   indent(): string {
     return this._indentController.indent();
   }