impr: Ignore console.logs in vertex shaders

apps/typegpu-docs/src/examples/tests/log-test/index.ts

@@ -215,16 +215,56 @@ export const controls = {
   },
   'Render pipeline': {
     onButtonClick: () => {
-      const context = canvas.getContext('webgpu') as GPUCanvasContext;
+      const pipeline = root['~unstable']
+        .withVertex(mainVertex)
+        .withFragment(mainFragment, { format: presentationFormat })
+        .createPipeline();
 
-      context.configure({
-        device: root.device,
-        format: presentationFormat,
-        alphaMode: 'premultiplied',
+      pipeline
+        .withColorAttachment({
+          view: context.getCurrentTexture().createView(),
+          clearValue: [0, 0, 0, 0],
+          loadOp: 'clear',
+          storeOp: 'store',
+        })
+        .draw(3);
+    },
+  },
+  'Render pipeline helper': {
+    onButtonClick: () => {
+      const helper = tgpu.fn([])(() => {
+        'use gpu';
+        console.log('A log made from a helper used during both stages.');
+      });
+
+      const mainVertex = tgpu['~unstable'].vertexFn({
+        in: { vertexIndex: d.builtin.vertexIndex },
+        out: { pos: d.builtin.position },
+      })((input) => {
+        console.log('A log made from vertex stage.');
+        helper();
+
+        const positions = [
+          d.vec2f(0, 0.5),
+          d.vec2f(-0.5, -0.5),
+          d.vec2f(0.5, -0.5),
+        ];
+
+        return { pos: d.vec4f(positions[input.vertexIndex], 0, 1) };
+      });
+
+      const mainFragment = tgpu['~unstable'].fragmentFn({
+        in: { pos: d.builtin.position },
+        out: d.vec4f,
+      })(() => {
+        console.log('A log made from fragment stage.');
+        helper();
+
+        return d.vec4f(0.769, 0.392, 1.0, 1);
       });
 
       const pipeline = root['~unstable']
-        .withVertex(mainVertex, {})
+        .withVertex(mainVertex)
         .withFragment(mainFragment, { format: presentationFormat })
         .createPipeline();
 

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

@@ -9,6 +9,7 @@ import {
 } from '../../shared/meta.ts';
 import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
+import { shaderStageSlot } from '../slot/internalSlots.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
 import type { Implementation, InferIO, IORecord } from './fnTypes.ts';
 import { createIoSchema, type IOLayoutToSchema } from './ioSchema.ts';
@@ -27,7 +28,7 @@ type TgpuComputeFnShellHeader<
   readonly argTypes: [IOLayoutToSchema<ComputeIn>] | [];
   readonly returnType: Void;
   readonly workgroupSize: [number, number, number];
-  readonly isEntry: true;
+  readonly entryPoint: 'compute';
 };
 
 /**
@@ -107,7 +108,7 @@ export function computeFn<
       options.workgroupSize[1] ?? 1,
       options.workgroupSize[2] ?? 1,
     ],
-    isEntry: true,
+    entryPoint: 'compute',
   };
 
   const call = (
@@ -123,6 +124,12 @@ export function computeFn<
   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';
+}
+
 // --------------
 // Implementation
 // --------------
@@ -162,11 +169,12 @@ function createComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
     },
 
     [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
-      return core.resolve(
-        ctx,
-        shell.argTypes,
-        shell.returnType,
-      );
+      return ctx.withSlots([[shaderStageSlot, 'compute']], () =>
+        core.resolve(
+          ctx,
+          shell.argTypes,
+          shell.returnType,
+        ));
     },
 
     toString() {

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

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

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

@@ -22,6 +22,7 @@ import {
 import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { addReturnTypeToExternals } from '../resolve/externals.ts';
+import { shaderStageSlot } from '../slot/internalSlots.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
 import type {
   BaseIOData,
@@ -61,7 +62,7 @@ type TgpuFragmentFnShellHeader<
   readonly in: FragmentIn | undefined;
   readonly out: FragmentOut;
   readonly returnType: IOLayoutToSchema<FragmentOut>;
-  readonly isEntry: true;
+  readonly entryPoint: 'fragment';
 };
 
 /**
@@ -145,7 +146,7 @@ export function fragmentFn<
     in: options.in,
     out: options.out,
     returnType: createIoSchema(options.out),
-    isEntry: true,
+    entryPoint: 'fragment',
   };
 
   const call = (
@@ -159,6 +160,17 @@ export function fragmentFn<
   >;
 }
 
+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';
+}
+
 // --------------
 // Implementation
 // --------------
@@ -215,11 +227,12 @@ function createFragmentFn(
       }
       core.applyExternals({ Out: outputType });
 
-      return core.resolve(
-        ctx,
-        inputWithLocation ? [inputWithLocation] : [],
-        shell.returnType,
-      );
+      return ctx.withSlots([[shaderStageSlot, 'fragment']], () =>
+        core.resolve(
+          ctx,
+          inputWithLocation ? [inputWithLocation] : [],
+          shell.returnType,
+        ));
     },
 
     toString() {

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

@@ -18,6 +18,7 @@ import {
 } from '../../shared/meta.ts';
 import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
+import { shaderStageSlot } from '../slot/internalSlots.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
 import type {
   BaseIOData,
@@ -52,7 +53,7 @@ type TgpuVertexFnShellHeader<
   readonly in: VertexIn | undefined;
   readonly out: VertexOut;
   readonly argTypes: [IOLayoutToSchema<VertexIn>] | [];
-  readonly isEntry: true;
+  readonly entryPoint: 'vertex';
 };
 
 /**
@@ -134,7 +135,7 @@ export function vertexFn<
     argTypes: options.in && Object.keys(options.in).length !== 0
       ? [createIoSchema(options.in)]
       : [],
-    isEntry: true,
+    entryPoint: 'vertex',
   };
 
   const call = (
@@ -145,6 +146,16 @@ export function vertexFn<
   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';
+}
+
 // --------------
 // Implementation
 // --------------
@@ -195,11 +206,12 @@ function createVertexFn(
         core.applyExternals({ Out: outputWithLocation });
       }
 
-      return core.resolve(
-        ctx,
-        shell.argTypes,
-        outputWithLocation,
-      );
+      return ctx.withSlots([[shaderStageSlot, 'vertex']], () =>
+        core.resolve(
+          ctx,
+          shell.argTypes,
+          outputWithLocation,
+        ));
     },
 
     toString() {

packages/typegpu/src/core/slot/internalSlots.ts

@@ -0,0 +1,4 @@
+import type { ShaderStage } from '../../types.ts';
+import { slot } from './slot.ts';
+
+export const shaderStageSlot = slot<ShaderStage>(undefined);

packages/typegpu/src/index.ts

@@ -108,6 +108,9 @@ export { isUsableAsStorage } from './extension.ts';
 export { isUsableAsUniform } from './core/buffer/bufferUsage.ts';
 export { isBufferShorthand } from './core/buffer/bufferShorthand.ts';
 export { isTgpuFn } from './core/function/tgpuFn.ts';
+export { isTgpuFragmentFn } from './core/function/tgpuFragmentFn.ts';
+export { isTgpuVertexFn } from './core/function/tgpuVertexFn.ts';
+export { isTgpuComputeFn } from './core/function/tgpuComputeFn.ts';
 export { isVariable } from './core/variable/tgpuVariable.ts';
 
 // types

packages/typegpu/src/tgsl/consoleLog/logGenerator.ts

@@ -1,6 +1,7 @@
 import type { TgpuMutable } from '../../core/buffer/bufferShorthand.ts';
 import { stitch } from '../../core/resolve/stitch.ts';
 import type { TgpuRoot } from '../../core/root/rootTypes.ts';
+import { shaderStageSlot } from '../../core/slot/internalSlots.ts';
 import { arrayOf } from '../../data/array.ts';
 import { atomic } from '../../data/atomic.ts';
 import { UnknownData } from '../../data/dataTypes.ts';
@@ -89,8 +90,13 @@ export class LogGeneratorImpl implements LogGenerator {
     op: string,
     args: Snippet[],
   ): Snippet {
+    if (shaderStageSlot.$ === 'vertex') {
+      console.warn(`'console.${op}' is not supported in vertex shaders.`);
+      return fallbackSnippet;
+    }
+
     if (!supportedLogOps.includes(op as SupportedLogOps)) {
-      console.warn(`Unsupported log method '${op}' was used in TGSL.`);
+      console.warn(`Unsupported log method '${op}'.`);
       return fallbackSnippet;
     }
 

packages/typegpu/src/types.ts

@@ -233,6 +233,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