feat: Include trace sources in errors

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

@@ -69,10 +69,9 @@ type AutoFragmentFnImpl = (
  */
 export class AutoFragmentFn implements SelfResolvable {
   // Prototype properties
-  declare [$internal]: true;
   declare resourceType: 'auto-fragment-fn';
 
-  #core: FnCore;
+  readonly [$internal]: { fnCore: FnCore };
   autoIn: AutoStruct;
   autoOut: AutoStruct;
 
@@ -85,7 +84,7 @@ export class AutoFragmentFn implements SelfResolvable {
     if (!getName(impl)) {
       setName(impl, 'fragmentFn');
     }
-    this.#core = createFnCore(impl, 'fragment');
+    this[$internal] = { fnCore: createFnCore(impl, 'fragment') };
     this.autoIn = new AutoStruct(
       { ...builtinFragmentIn, ...varyings },
       undefined,
@@ -103,12 +102,11 @@ export class AutoFragmentFn implements SelfResolvable {
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     return ctx.withSlots([[shaderStageSlot, 'fragment']], () =>
-      this.#core.resolve(ctx, [this.autoIn], this.autoOut),
+      this[$internal].fnCore.resolve(ctx, [this.autoIn], this.autoOut),
     );
   }
 }
 
-AutoFragmentFn.prototype[$internal] = true;
 AutoFragmentFn.prototype.resourceType = 'auto-fragment-fn';
 
 type AutoVertexFnImpl = (
@@ -120,10 +118,9 @@ type AutoVertexFnImpl = (
  */
 export class AutoVertexFn implements SelfResolvable {
   // Prototype properties
-  declare [$internal]: true;
   declare resourceType: 'auto-vertex-fn';
 
-  #core: FnCore;
+  readonly [$internal]: { fnCore: FnCore };
   autoIn: AutoStruct;
   autoOut: AutoStruct;
 
@@ -136,7 +133,7 @@ export class AutoVertexFn implements SelfResolvable {
     if (!getName(impl)) {
       setName(impl, 'vertexFn');
     }
-    this.#core = createFnCore(impl, 'vertex');
+    this[$internal] = { fnCore: createFnCore(impl, 'vertex') };
     this.autoIn = new AutoStruct({ ...builtinVertexIn, ...attribs }, undefined);
     setName(this.autoIn, 'VertexIn');
     this.autoOut = new AutoStruct(
@@ -154,10 +151,9 @@ export class AutoVertexFn implements SelfResolvable {
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
     return ctx.withSlots([[shaderStageSlot, 'vertex']], () =>
-      this.#core.resolve(ctx, [this.autoIn], this.autoOut),
+      this[$internal].fnCore.resolve(ctx, [this.autoIn], this.autoOut),
     );
   }
 }
 
-AutoVertexFn.prototype[$internal] = true;
 AutoVertexFn.prototype.resourceType = 'auto-vertex-fn';

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

@@ -6,7 +6,7 @@ import { type BaseData, isWgslData, isWgslStruct, Void } from '../../data/wgslTy
 import { validateIdentifier } from '../../nameUtils.ts';
 import { getFunctionMetadata, getName } from '../../shared/meta.ts';
 import { $getNameForward } from '../../shared/symbols.ts';
-import type { ResolutionCtx, ShaderStage } from '../../types.ts';
+import type { Source, ResolutionCtx, ShaderStage } from '../../types.ts';
 import {
   type ExternalMap,
   replaceExternalsInWgsl,
@@ -56,6 +56,8 @@ export interface FnCore {
      */
     entryInput?: SeparatedEntryArgs,
   ): ResolvedSnippet;
+  /** Points to the implementation's body. Filled in during resolution. */
+  source: Source | undefined;
 }
 
 export function createFnCore(
@@ -90,6 +92,8 @@ export function createFnCore(
       }
     },
 
+    source: undefined as undefined | Source,
+
     resolve(
       ctx: ResolutionCtx,
       argTypes: BaseData[],
@@ -225,6 +229,13 @@ export function createFnCore(
       const [body, sourceMap] = stripSourceMap(ast.body) as [Block, SourceMap];
       if (pluginData.filename !== undefined) {
         ctx.addSourceMap(sourceMap, pluginData.filename);
+        const maybeSource = sourceMap.get(body);
+        if (maybeSource) {
+          this.source = {
+            file: pluginData.filename,
+            location: maybeSource,
+          };
+        }
       }
 
       // If an entrypoint implementation has a second argument, it represents the output schema.

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

@@ -3,7 +3,7 @@ import type { BaseData } from '../../data/wgslTypes.ts';
 import { getName } from '../../shared/meta.ts';
 import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
-import { createFnCore } from './fnCore.ts';
+import { createFnCore, type FnCore } from './fnCore.ts';
 
 /**
  * Shell-less functions are possible because we can infer the signature based solely on the context
@@ -24,6 +24,7 @@ import { createFnCore } from './fnCore.ts';
 export interface ShelllessImpl extends SelfResolvable {
   readonly resourceType: 'shellless-impl';
   readonly argTypes: BaseData[];
+  readonly [$internal]: { fnCore: FnCore };
   readonly [$getNameForward]: unknown;
 }
 
@@ -38,7 +39,7 @@ export function createShelllessImpl(
   const core = createFnCore(implementation, 'normal');
 
   return {
-    [$internal]: true,
+    [$internal]: { fnCore: core },
     [$getNameForward]: core,
     resourceType: 'shellless-impl' as const,
     argTypes,

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

@@ -47,7 +47,7 @@ export interface TgpuComputeFn<
   // oxlint-disable-next-line typescript/no-explicit-any -- to allow assigning any compute fn to TgpuComputeFn (non-generic) type
   ComputeIn extends IORecord<AnyComputeBuiltin> = any,
 > extends TgpuNamable {
-  readonly [$internal]: true;
+  readonly [$internal]: { fnCore: FnCore };
   readonly shell: TgpuComputeFnShellHeader<ComputeIn>;
 
   $uses(dependencyMap: Record<string, unknown>): this;
@@ -110,7 +110,6 @@ function createComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
 ): TgpuComputeFn<ComputeIn> {
   type This = TgpuComputeFn<ComputeIn> &
     SelfResolvable & {
-      [$internal]: true;
       [$getNameForward]: FnCore;
     };
 
@@ -124,7 +123,7 @@ function createComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
       return this;
     },
 
-    [$internal]: true,
+    [$internal]: { fnCore: core },
     [$getNameForward]: core,
     $name(newLabel: string): This {
       setName(this, newLabel);

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

@@ -69,6 +69,7 @@ export type TgpuFnShell<Args extends BaseData[], Return extends BaseData> = Tgpu
 interface TgpuFnBase<ImplSchema extends AnyFn> extends TgpuNamable, Withable<TgpuFn<ImplSchema>> {
   [$internal]: {
     implementation: Implementation<ImplSchema>;
+    fnCore: FnCore;
   };
   readonly resourceType: 'function';
   readonly shell: TgpuFnShellHeader<
@@ -212,7 +213,7 @@ function createFn<ImplSchema extends AnyFn>(
   const fnBase = {
     shell,
     resourceType: 'function' as const,
-    [$internal]: { implementation },
+    [$internal]: { implementation, fnCore: core },
 
     $uses(newExternals: Record<string, unknown>) {
       core.setExternals('userProvided', newExternals);
@@ -303,7 +304,7 @@ function createBoundFunction<ImplSchema extends AnyFn>(
   const fnBase: This = {
     resourceType: 'function',
     shell: innerFn.shell,
-    [$internal]: { implementation: innerFn[$internal].implementation },
+    [$internal]: { ...innerFn[$internal] },
     [$providing]: { inner: innerFn, pairs },
 
     $uses(newExternals) {

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

@@ -110,7 +110,7 @@ export interface TgpuFragmentFn<
   in Varying extends TgpuFragmentFn.In = Record<string, never>,
   out Output extends TgpuFragmentFn.Out = TgpuFragmentFn.Out,
 > extends TgpuNamable {
-  readonly [$internal]: true;
+  readonly [$internal]: { fnCore: FnCore };
   readonly shell: TgpuFragmentFnShellHeader<Varying, Output>;
   readonly outputType: IOLayoutToSchema<Output>;
 
@@ -182,7 +182,6 @@ function createFragmentFn(
 ): TgpuFragmentFn {
   type This = TgpuFragmentFn<TgpuFragmentFn.In> &
     SelfResolvable & {
-      [$internal]: true;
       [$getNameForward]: FnCore;
     };
 
@@ -201,7 +200,7 @@ function createFragmentFn(
       return this;
     },
 
-    [$internal]: true,
+    [$internal]: { fnCore: core },
     [$getNameForward]: core,
     $name(newLabel: string): This {
       setName(this, newLabel);

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

@@ -81,7 +81,7 @@ export interface TgpuVertexFn<
   in VertexIn extends TgpuVertexFn.In = Record<string, never>,
   out VertexOut extends TgpuVertexFn.Out = TgpuVertexFn.Out,
 > extends TgpuNamable {
-  readonly [$internal]: true;
+  readonly [$internal]: { fnCore: FnCore };
   readonly shell: TgpuVertexFnShellHeader<VertexIn, VertexOut>;
   $uses(dependencyMap: Record<string, unknown>): this;
 }
@@ -159,7 +159,6 @@ function createVertexFn(
 ): TgpuVertexFn<VertexInConstrained, VertexOutConstrained> {
   type This = TgpuVertexFn<VertexInConstrained, VertexOutConstrained> &
     SelfResolvable & {
-      [$internal]: true;
       [$getNameForward]: FnCore;
     };
 
@@ -174,7 +173,7 @@ function createVertexFn(
       return this;
     },
 
-    [$internal]: true,
+    [$internal]: { fnCore: core },
     [$getNameForward]: core,
     $name(newLabel: string): This {
       setName(this, newLabel);

packages/typegpu/src/errors.ts

@@ -7,7 +7,9 @@ import { DEV, INCLUDE_SOURCE, TEST } from './shared/env.ts';
 import type { TgpuBindGroupLayout } from './tgpuBindGroupLayout.ts';
 import type * as tinyest from 'tinyest';
 import { stringifyNode } from './shared/tseynit.ts';
-import type { ErrorSource } from './types.ts';
+import type { Source } from './types.ts';
+import { $internal } from './shared/symbols.ts';
+import type { FnCore } from './core/function/fnCore.ts';
 
 const prefix = 'Invariant failed';
 
@@ -40,6 +42,35 @@ export function invariant(
   throw new Error(value);
 }
 
+function getFnSource(item: unknown): Source | undefined {
+  const typed = item as { [$internal]?: { fnCore?: FnCore } } | undefined;
+  return typed?.[$internal]?.fnCore?.source;
+}
+
+function stringifySource(source: Source): string {
+  return `${source.file}:${source.location[0]}:${source.location[1]}`;
+}
+
+function constructMessage(cause: unknown, trace: unknown[], source?: Source) {
+  const causeString =
+    cause && typeof cause === 'object' && 'message' in cause ? cause.message : cause;
+
+  const entries = trace.length > 20 ? [...trace.slice(0, 11), '...', ...trace.slice(-10)] : trace;
+  const traceString = entries
+    .map((ancestor, idx) => {
+      const base = `- ${hasTinyestMetadata(ancestor) ? `fn*:${getName(ancestor)}` : ancestor}`;
+      if (idx < entries.length - 1) {
+        const maybeSource = getFnSource(ancestor);
+        return `${base}${INCLUDE_SOURCE && maybeSource !== undefined ? ` (defined at ${stringifySource(maybeSource)})` : ''}`;
+      } else {
+        return `${base}${INCLUDE_SOURCE && source !== undefined ? ` (at ${stringifySource(source)}):\n` : ': '}`;
+      }
+    })
+    .join('\n');
+
+  return `Resolution of the following tree failed:\n${traceString}${causeString}`;
+}
+
 /**
  * An error that happens during resolution of WGSL code.
  * Contains a trace of all ancestor resolvables in
@@ -50,27 +81,10 @@ export function invariant(
 export class ResolutionError extends Error {
   readonly cause: unknown;
   readonly trace: unknown[];
-  readonly location: ErrorSource | undefined;
-
-  constructor(cause: unknown, trace: unknown[], source?: ErrorSource) {
-    let entries = trace.map(
-      (ancestor) => `- ${hasTinyestMetadata(ancestor) ? `fn*:${getName(ancestor)}` : ancestor}`,
-    );
-
-    // Showing only the root and leaf nodes.
-    if (entries.length > 20) {
-      entries = [...entries.slice(0, 11), '...', ...entries.slice(-10)];
-    }
-
-    const sourceString = `${entries.join('\n')}${
-      INCLUDE_SOURCE && source !== undefined
-        ? ` (at ${source.file}:${source.location[0]}:${source.location[1]}):\n`
-        : `: `
-    }`;
-    const causeString =
-      cause && typeof cause === 'object' && 'message' in cause ? cause.message : cause;
+  readonly location: Source | undefined;
 
-    super(`Resolution of the following tree failed:\n${sourceString}${causeString}`);
+  constructor(cause: unknown, trace: unknown[], source?: Source) {
+    super(constructMessage(cause, trace, source));
 
     this.cause = cause;
     this.trace = trace;

packages/typegpu/src/types.ts

@@ -268,8 +268,8 @@ export type ExecState = NormalState | CodegenState | SimulationState;
 
 type MapKey = SourceMap extends Map<infer K, unknown> ? K : never;
 type MapValue = SourceMap extends Map<unknown, infer V> ? V : never;
-export type ErrorSource = { file: string | undefined; location: MapValue };
-export type CombinedSourceMap = Map<MapKey, ErrorSource>;
+export type Source = { file: string | undefined; location: MapValue };
+export type CombinedSourceMap = Map<MapKey, Source>;
 
 /**
  * Passed into each resolvable item. All items in a tree share a resolution ctx,