fix: Fill missing slots in shellless entry functions (#2255)

Author: aleksanderkatan

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

@@ -1,4 +1,4 @@
-import tgpu, { d, std } from 'typegpu';
+import tgpu, { d, std, type AutoVertexIn } from 'typegpu';
 import { defineControls } from '../../common/defineControls.ts';
 
 const root = await tgpu.init({
@@ -249,6 +249,32 @@ export const controls = defineControls({
       pipeline.withIndexBuffer(indexBuffer).withColorAttachment({ view: context }).drawIndexed(3);
     },
   },
+  'Shellless entry': {
+    onButtonClick: () => {
+      const myLog = (n: number) => {
+        'use gpu';
+        console.log(n);
+      };
+
+      const vs = ({ $vertexIndex }: AutoVertexIn<{}>) => {
+        'use gpu';
+        const positions = [d.vec2f(0, 0.5), d.vec2f(-0.5, -0.5), d.vec2f(0.5, -0.5)];
+        myLog(6);
+        return { $position: d.vec4f(positions[$vertexIndex], 0, 1) };
+      };
+      const fs = () => {
+        'use gpu';
+        myLog(7);
+        return d.vec4f();
+      };
+
+      const pipeline = root.createRenderPipeline({
+        vertex: vs,
+        fragment: fs,
+      });
+      pipeline.withColorAttachment({ view: context }).draw(3);
+    },
+  },
   'Too many logs': {
     onButtonClick: () =>
       root

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

@@ -7,6 +7,7 @@ import { getName, setName } from '../../shared/meta.ts';
 import type { InferGPU, InferGPURecord, InferRecord } from '../../shared/repr.ts';
 import { $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 } from './fnTypes.ts';
 
@@ -85,7 +86,9 @@ export class AutoFragmentFn implements SelfResolvable {
   }
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
-    return this.#core.resolve(ctx, [this.#autoIn], this.#autoOut);
+    return ctx.withSlots([[shaderStageSlot, 'fragment']], () =>
+      this.#core.resolve(ctx, [this.#autoIn], this.#autoOut),
+    );
   }
 }
 
@@ -129,7 +132,9 @@ export class AutoVertexFn implements SelfResolvable {
   }
 
   [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
-    return this.#core.resolve(ctx, [this.#autoIn], this.#autoOut);
+    return ctx.withSlots([[shaderStageSlot, 'vertex']], () =>
+      this.#core.resolve(ctx, [this.#autoIn], this.#autoOut),
+    );
   }
 }
 

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

@@ -223,6 +223,92 @@ describe('wgslGenerator with console.log', () => {
     `);
   });
 
+  it('Works for shellless entry functions', ({ root }) => {
+    const myLog = (n: number) => {
+      'use gpu';
+      console.log(n);
+    };
+
+    const vs = () => {
+      'use gpu';
+      myLog(6);
+      return { pos: d.vec4f() };
+    };
+    const fs = () => {
+      'use gpu';
+      myLog(7);
+      return d.vec4f();
+    };
+
+    const pipeline = root.createRenderPipeline({
+      vertex: vs,
+      fragment: fs,
+      targets: { format: 'rg8unorm' },
+    });
+
+    expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(`
+      "fn myLog(n: i32) {
+        /* console.log() */;
+      }
+
+      struct VertexOut {
+        @location(0) pos: vec4f,
+      }
+
+      @vertex fn vs() -> VertexOut {
+        myLog(6i);
+        return VertexOut(vec4f());
+      }
+
+      @group(0) @binding(0) var<storage, read_write> indexBuffer: atomic<u32>;
+
+      struct SerializedLogData {
+        id: u32,
+        serializedData: array<u32, 63>,
+      }
+
+      @group(0) @binding(1) var<storage, read_write> dataBuffer: array<SerializedLogData, 64>;
+
+      var<private> dataBlockIndex: u32;
+
+      var<private> dataByteIndex: u32;
+
+      fn nextByteIndex() -> u32{
+        let i = dataByteIndex;
+        dataByteIndex = dataByteIndex + 1u;
+        return i;
+      }
+
+      fn serializeI32(n: i32) {
+        dataBuffer[dataBlockIndex].serializedData[nextByteIndex()] = bitcast<u32>(n);
+      }
+
+      fn log1serializer(_arg_0: i32) {
+        serializeI32(_arg_0);
+      }
+
+      fn log1(_arg_0: i32) {
+        dataBlockIndex = atomicAdd(&indexBuffer, 1);
+        if (dataBlockIndex >= 64) {
+          return;
+        }
+        dataBuffer[dataBlockIndex].id = 1;
+        dataByteIndex = 0;
+
+        log1serializer(_arg_0);
+      }
+
+      fn myLog_1(n: i32) {
+        log1(n);
+      }
+
+      @fragment fn fs() -> @location(0) vec4f {
+        myLog_1(7i);
+        return vec4f();
+      }"
+    `);
+  });
+
   it('Parses a single console.log in a compute pipeline', ({ root }) => {
     const fn = tgpu.computeFn({
       workgroupSize: [1],