feat: Expected type stack in WgslGenerator

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();
   }

packages/typegpu/src/shared/generators.ts

@@ -1,10 +1,3 @@
-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 './meta.ts';
-import { $internal } from './symbols.ts';
-
 /**
  * Yields values in the sequence 0,1,2..∞ except for the ones in the `excluded` set.
  */
@@ -21,23 +14,3 @@ export function* naturalsExcept(
     next++;
   }
 }
-
-export function createDualImpl<T extends (...args: never[]) => unknown>(
-  jsImpl: T,
-  gpuImpl: (...args: MapValueToSnippet<Parameters<T>>) => Snippet,
-  name: string,
-  argTypes?: FnArgsConversionHint,
-): 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, argTypes };
-
-  return impl as TgpuDualFn<T>;
-}

packages/typegpu/src/std/array.ts

@@ -3,7 +3,7 @@ import { abstractInt, u32 } from '../data/numeric.ts';
 import { ptrFn } from '../data/ptr.ts';
 import type { AnyWgslData } from '../data/wgslTypes.ts';
 import { isPtr, isWgslArray } from '../data/wgslTypes.ts';
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 
 export const arrayLength = createDualImpl(
   // CPU implementation

packages/typegpu/src/std/atomic.ts

@@ -7,7 +7,7 @@ import {
   isWgslData,
   Void,
 } from '../data/wgslTypes.ts';
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 type AnyAtomic = atomicI32 | atomicU32;
 
 export const workgroupBarrier = createDualImpl(

packages/typegpu/src/std/boolean.ts

@@ -14,7 +14,7 @@ import {
   type v3b,
   type v4b,
 } from '../data/wgslTypes.ts';
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 import { isSnippetNumeric, sub } from './numeric.ts';
 
 function correspondingBooleanVectorSchema(value: Snippet) {

packages/typegpu/src/std/discard.ts

@@ -1,6 +1,6 @@
 import { snip } from '../data/snippet.ts';
 import { Void } from '../data/wgslTypes.ts';
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 
 export const discard = createDualImpl(
   // CPU

packages/typegpu/src/std/matrix.ts

@@ -1,7 +1,7 @@
 import type { Snippet } from '../data/snippet.ts';
 import { mat4x4f } from '../data/matrix.ts';
 import type { m4x4f, v3f } from '../data/wgslTypes.ts';
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 import { mul } from './numeric.ts';
 
 /**

packages/typegpu/src/std/numeric.ts

@@ -29,7 +29,7 @@ import {
   type v4i,
   type vBaseForMat,
 } from '../data/wgslTypes.ts';
-import { createDualImpl } from '../shared/generators.ts';
+import { createDualImpl } from '../core/function/dualImpl.ts';
 import { $internal } from '../shared/symbols.ts';
 
 type NumVec = AnyNumericVecInstance;
@@ -96,7 +96,7 @@ export const add = createDualImpl(
       `(${lhs.value} + ${rhs.value})`,
       isSnippetNumeric(lhs) ? rhs.dataType : lhs.dataType,
     ),
-  'coerce',
+  'unify',
 );
 
 function cpuSub(lhs: number, rhs: number): number; // default subtraction
@@ -126,7 +126,7 @@ export const sub = createDualImpl(
       isSnippetNumeric(lhs) ? rhs.dataType : lhs.dataType,
     ),
   'sub',
-  'coerce',
+  'unify',
 );
 
 function cpuMul(lhs: number, rhs: number): number; // default multiplication
@@ -510,7 +510,7 @@ export const max = createDualImpl(
   // GPU implementation
   (a, b) => snip(`max(${a.value}, ${b.value})`, a.dataType),
   'max',
-  'coerce',
+  'unify',
 );
 
 /**
@@ -528,7 +528,7 @@ export const min = createDualImpl(
   // GPU implementation
   (a, b) => snip(`min(${a.value}, ${b.value})`, a.dataType),
   'min',
-  'coerce',
+  'unify',
 );
 
 export const sign = createDualImpl(