impr: Better error when logging from vertex shader

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

@@ -27,7 +27,7 @@ type TgpuComputeFnShellHeader<
   readonly argTypes: [IOLayoutToSchema<ComputeIn>] | [];
   readonly returnType: Void;
   readonly workgroupSize: [number, number, number];
-  readonly isEntry: true;
+  readonly entryPoint: 'compute';
 };
 
 /**
@@ -54,22 +54,7 @@ export type TgpuComputeFnShell<
   & ((
     strings: TemplateStringsArray,
     ...values: unknown[]
-  ) => TgpuComputeFn<ComputeIn>)
-  & {
-    /**
-     * @deprecated Invoke the shell as a function instead.
-     */
-    does:
-      & ((
-        implementation: (input: InferIO<ComputeIn>) => undefined,
-      ) => TgpuComputeFn<ComputeIn>)
-      & /**
-       * @param implementation
-       *   Raw WGSL function implementation with header and body
-       *   without `fn` keyword and function name
-       *   e.g. `"(x: f32) -> f32 { return x; }"`;
-       */ ((implementation: string) => TgpuComputeFn<ComputeIn>);
-  };
+  ) => TgpuComputeFn<ComputeIn>);
 
 export interface TgpuComputeFn<
   // biome-ignore lint/suspicious/noExplicitAny: to allow assigning any compute fn to TgpuComputeFn (non-generic) type
@@ -122,7 +107,7 @@ export function computeFn<
       options.workgroupSize[1] ?? 1,
       options.workgroupSize[2] ?? 1,
     ],
-    isEntry: true,
+    entryPoint: 'compute',
   };
 
   const call = (
@@ -135,9 +120,13 @@ export function computeFn<
       stripTemplate(arg, ...values),
     );
 
-  return Object.assign(Object.assign(call, shell), {
-    does: call,
-  }) as TgpuComputeFnShell<ComputeIn>;
+  return Object.assign(call, shell);
+}
+
+export function isTgpuComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
+  value: unknown | TgpuComputeFn<ComputeIn>,
+): value is TgpuComputeFn<ComputeIn> {
+  return (value as TgpuComputeFn<ComputeIn>)?.shell?.entryPoint === 'compute';
 }
 
 // --------------

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

@@ -56,7 +56,6 @@ type TgpuFnShellHeader<
   readonly [$internal]: true;
   readonly argTypes: Args;
   readonly returnType: Return;
-  readonly isEntry: false;
 };
 
 /**
@@ -125,7 +124,6 @@ export function fn<
     [$internal]: true,
     argTypes,
     returnType: returnType ?? Void as Return,
-    isEntry: false,
   };
 
   const call = (

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

@@ -61,7 +61,7 @@ type TgpuFragmentFnShellHeader<
   readonly in: FragmentIn | undefined;
   readonly out: FragmentOut;
   readonly returnType: IOLayoutToSchema<FragmentOut>;
-  readonly isEntry: true;
+  readonly entryPoint: 'fragment';
 };
 
 /**
@@ -93,24 +93,7 @@ export type TgpuFragmentFnShell<
   & ((
     strings: TemplateStringsArray,
     ...values: unknown[]
-  ) => TgpuFragmentFn<OmitBuiltins<FragmentIn>, FragmentOut>)
-  & {
-    /**
-     * @deprecated Invoke the shell as a function instead.
-     */
-    does:
-      & ((
-        implementation: (input: InferIO<FragmentIn>) => InferIO<FragmentOut>,
-      ) => TgpuFragmentFn<OmitBuiltins<FragmentIn>, FragmentOut>)
-      & /**
-       * @param implementation
-       *   Raw WGSL function implementation with header and body
-       *   without `fn` keyword and function name
-       *   e.g. `"(x: f32) -> f32 { return x; }"`;
-       */ ((
-        implementation: string,
-      ) => TgpuFragmentFn<OmitBuiltins<FragmentIn>, FragmentOut>);
-  };
+  ) => TgpuFragmentFn<OmitBuiltins<FragmentIn>, FragmentOut>);
 
 export interface TgpuFragmentFn<
   Varying extends FragmentInConstrained = FragmentInConstrained,
@@ -162,17 +145,29 @@ export function fragmentFn<
     in: options.in,
     out: options.out,
     returnType: createIoSchema(options.out),
-    isEntry: true,
+    entryPoint: 'fragment',
   };
 
   const call = (
     arg: Implementation | TemplateStringsArray,
     ...values: unknown[]
   ) => createFragmentFn(shell, stripTemplate(arg, ...values));
 
-  return Object.assign(Object.assign(call, shell), {
-    does: call,
-  }) as TgpuFragmentFnShell<FragmentIn, FragmentOut>;
+  return Object.assign(call, shell) as TgpuFragmentFnShell<
+    FragmentIn,
+    FragmentOut
+  >;
+}
+
+export function isTgpuFragmentFn<
+  FragmentIn extends FragmentInConstrained,
+  FragmentOut extends FragmentOutConstrained,
+>(
+  value: unknown | TgpuFragmentFn<FragmentIn, FragmentOut>,
+): value is TgpuFragmentFn<FragmentIn, FragmentOut> {
+  return (value as TgpuFragmentFn<FragmentIn, FragmentOut>)?.shell
+    ?.entryPoint ===
+    'fragment';
 }
 
 // --------------

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

@@ -52,7 +52,7 @@ type TgpuVertexFnShellHeader<
   readonly in: VertexIn | undefined;
   readonly out: VertexOut;
   readonly argTypes: [IOLayoutToSchema<VertexIn>] | [];
-  readonly isEntry: true;
+  readonly entryPoint: 'vertex';
 };
 
 /**
@@ -77,19 +77,7 @@ export type TgpuVertexFnShell<
   & ((
     strings: TemplateStringsArray,
     ...values: unknown[]
-  ) => TgpuVertexFn<OmitBuiltins<VertexIn>, OmitBuiltins<VertexOut>>)
-  & {
-    /**
-     * @deprecated Invoke the shell as a function instead.
-     */
-    does:
-      & ((
-        implementation: (input: InferIO<VertexIn>) => InferIO<VertexOut>,
-      ) => TgpuVertexFn<OmitBuiltins<VertexIn>, OmitBuiltins<VertexOut>>)
-      & ((
-        implementation: string,
-      ) => TgpuVertexFn<OmitBuiltins<VertexIn>, OmitBuiltins<VertexOut>>);
-  };
+  ) => TgpuVertexFn<OmitBuiltins<VertexIn>, OmitBuiltins<VertexOut>>);
 
 export interface TgpuVertexFn<
   VertexIn extends VertexInConstrained = VertexInConstrained,
@@ -146,17 +134,25 @@ export function vertexFn<
     argTypes: options.in && Object.keys(options.in).length !== 0
       ? [createIoSchema(options.in)]
       : [],
-    isEntry: true,
+    entryPoint: 'vertex',
   };
 
   const call = (
     arg: Implementation | TemplateStringsArray,
     ...values: unknown[]
   ) => createVertexFn(shell, stripTemplate(arg, ...values));
 
-  return Object.assign(Object.assign(call, shell), {
-    does: call,
-  }) as TgpuVertexFnShell<VertexIn, VertexOut>;
+  return Object.assign(call, shell) as TgpuVertexFnShell<VertexIn, VertexOut>;
+}
+
+export function isTgpuVertexFn<
+  VertexIn extends VertexInConstrained,
+  VertexOut extends VertexOutConstrained,
+>(
+  value: unknown | TgpuVertexFn<VertexIn, VertexOut>,
+): value is TgpuVertexFn<VertexIn, VertexOut> {
+  return (value as TgpuVertexFn<VertexIn, VertexOut>)?.shell?.entryPoint ===
+    'vertex';
 }
 
 // --------------

packages/typegpu/src/resolutionCtx.ts

@@ -65,12 +65,16 @@ import type {
   ItemLayer,
   ItemStateStack,
   ResolutionCtx,
+  ShaderStage,
   TgpuShaderStage,
   Wgsl,
 } from './types.ts';
 import { CodegenState, isSelfResolvable, NormalState } from './types.ts';
 import type { WgslExtension } from './wgslExtensions.ts';
 import { hasTinyestMetadata } from './shared/meta.ts';
+import { isTgpuComputeFn } from './core/function/tgpuComputeFn.ts';
+import { isTgpuVertexFn } from './core/function/tgpuVertexFn.ts';
+import { isTgpuFragmentFn } from './core/function/tgpuFragmentFn.ts';
 
 /**
  * Inserted into bind group entry definitions that belong
@@ -334,6 +338,12 @@ export class ResolutionCtxImpl implements ResolutionCtx {
   readonly #namespace: NamespaceInternal;
   readonly #shaderGenerator: ShaderGenerator;
 
+  /**
+   * Holds info about the currently resolved shader stage (if there is any).
+   * Note that if a function is used both in vertex and fragment stage,
+   * then it will only go through the process during the vertex stage.
+   */
+  private _currentStage: ShaderStage;
   private readonly _indentController = new IndentController();
   private readonly _itemStateStack = new ItemStateStackImpl();
   readonly #modeStack: ExecState[] = [];
@@ -444,6 +454,11 @@ export class ResolutionCtxImpl implements ResolutionCtx {
   }
 
   generateLog(op: string, args: Snippet[]): Snippet {
+    if (this._currentStage === 'vertex') {
+      throw new Error(
+        `'console.log' is not supported during vertex shader stage.`,
+      );
+    }
     return this.#logGenerator.generateLog(this, op, args);
   }
 
@@ -732,11 +747,24 @@ export class ResolutionCtxImpl implements ResolutionCtx {
     }
 
     if (isMarkedInternal(item) || hasTinyestMetadata(item)) {
+      // if we're resolving an entrypoint function, we want to update this._currentStage
+      const stage = getItemStage(item);
+      const resolutionAction = () => {
+        if (stage) {
+          this._currentStage = stage;
+        }
+        const result = this._getOrInstantiate(item);
+        if (stage) {
+          this._currentStage = undefined;
+        }
+        return result;
+      };
+
       // Top-level resolve
       if (this._itemStateStack.itemDepth === 0) {
         try {
           this.pushMode(new CodegenState());
-          const result = provideCtx(this, () => this._getOrInstantiate(item));
+          const result = provideCtx(this, resolutionAction);
           return snip(
             `${[...this._declarations].join('\n\n')}${result.value}`,
             Void,
@@ -747,7 +775,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
         }
       }
 
-      return this._getOrInstantiate(item);
+      return resolutionAction();
     }
 
     // This is a value that comes from the outside, maybe we can coerce it
@@ -972,3 +1000,15 @@ export function resolveFunctionHeader(
     } `
     : `(${argList}) `;
 }
+
+function getItemStage(item: unknown): ShaderStage {
+  if (isTgpuComputeFn(item)) {
+    return 'compute';
+  }
+  if (isTgpuVertexFn(item)) {
+    return 'vertex';
+  }
+  if (isTgpuFragmentFn(item)) {
+    return 'fragment';
+  }
+}

packages/typegpu/src/types.ts

@@ -217,6 +217,12 @@ export type ExecState =
   | CodegenState
   | SimulationState;
 
+export type ShaderStage =
+  | 'vertex'
+  | 'fragment'
+  | 'compute'
+  | undefined;
+
 /**
  * Passed into each resolvable item. All items in a tree share a resolution ctx,
  * but there can be layers added and removed from the item stack when going down

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

@@ -37,6 +37,74 @@ describe('wgslGenerator with console.log', () => {
     expect(consoleWarnSpy).toHaveBeenCalledTimes(1);
   });
 
+  it('Throws appropriate error when in a vertex shader', ({ root }) => {
+    const myLog = (n: number) => {
+      'use gpu';
+      console.log(n);
+    };
+
+    const vs = tgpu['~unstable'].vertexFn({ out: { pos: d.builtin.position } })(
+      () => {
+        myLog(5);
+        return { pos: d.vec4f() };
+      },
+    );
+    expect(() => tgpu.resolve([vs])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - vertexFn:vs
+      - fn*:myLog(i32)
+      - fn:consoleLog: 'console.log' is not supported during vertex shader stage.]
+    `);
+  });
+
+  it('Throws appropriate error when in a vertex shader during pipeline resolution', ({ root }) => {
+    const myLog = (n: number) => {
+      'use gpu';
+      console.log(n);
+    };
+
+    const vs = tgpu['~unstable'].vertexFn({ out: { pos: d.builtin.position } })(
+      () => {
+        myLog(5);
+        return { pos: d.vec4f() };
+      },
+    );
+    const fs = tgpu['~unstable'].fragmentFn({ out: d.vec4f })(() => {
+      return d.vec4f();
+    });
+
+    const pipeline = root['~unstable']
+      .withVertex(vs, {})
+      .withFragment(fs, { format: 'rg8unorm' })
+      .createPipeline();
+
+    expect(() => tgpu.resolve([pipeline])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - renderPipeline:pipeline
+      - renderPipelineCore
+      - vertexFn:vs
+      - fn*:myLog(i32)
+      - fn:consoleLog: 'console.log' is not supported during vertex shader stage.]
+    `);
+  });
+
+  it('Ignores console.log in a fragment shader resolved without a pipeline', () => {
+    const fs = tgpu['~unstable']
+      .fragmentFn({ out: d.vec4f })(() => {
+        console.log(d.u32(321));
+        return d.vec4f();
+      });
+
+    expect(tgpu.resolve([fs])).toMatchInlineSnapshot(`
+      "@fragment fn fs() -> @location(0) vec4f {
+        /* console.log() */;
+        return vec4f();
+      }"
+    `);
+  });
+
   it('Parses a single console.log in a render pipeline', ({ root }) => {
     const vs = tgpu['~unstable']
       .vertexFn({ out: { pos: d.builtin.position } })(() => {