Change vertexFn parameters to accept a struct each. Handle TGSL vertex shader resolution

apps/typegpu-docs/src/content/examples/rendering/box-raytracing/index.ts

@@ -185,7 +185,7 @@ const vertexFunction = tgpu['~unstable']
     { vertexIndex: d.builtin.vertexIndex },
     { outPos: d.builtin.position },
   )
-  .does(/* wgsl */ `(@builtin(vertex_index) vertexIndex: u32) -> VertexOutput {
+  .does(/* wgsl */ `(input: VertexInput) -> VertexOutput {
   var pos = array<vec2f, 6>(
     vec2<f32>( 1,  1),
     vec2<f32>( 1, -1),
@@ -196,7 +196,7 @@ const vertexFunction = tgpu['~unstable']
   );
 
   var output: VertexOutput;
-  output.outPos = vec4f(pos[vertexIndex], 0, 1);
+  output.outPos = vec4f(pos[input.vertexIndex], 0, 1);
   return output;
 }`)
   .$name('vertex_main');

apps/typegpu-docs/src/content/examples/simulation/boids-next/index.ts

@@ -42,11 +42,11 @@ const VertexOutput = {
 
 const mainVert = tgpu['~unstable']
   .vertexFn({ v: d.vec2f, center: d.vec2f, velocity: d.vec2f }, VertexOutput)
-  .does(/* wgsl */ `(@location(0) v: vec2f, @location(1) center: vec2f, @location(2) velocity: vec2f) -> VertexOutput {
-    let angle = getRotationFromVelocity(velocity);
-    let rotated = rotate(v, angle);
+  .does(/* wgsl */ `(input: VertexInput) -> VertexOutput {
+    let angle = getRotationFromVelocity(input.velocity);
+    let rotated = rotate(input.v, angle);
 
-    let pos = vec4(rotated + center, 0.0, 1.0);
+    let pos = vec4(rotated + input.center, 0.0, 1.0);
 
     let color = vec4(
         sin(angle + colorPalette.r) * 0.45 + 0.45,
@@ -223,7 +223,7 @@ const mainCompute = tgpu['~unstable']
     var alignmentCount = 0u;
     var cohesion = vec2(0.0, 0.0);
     var cohesionCount = 0u;
-    
+
     for (var i = 0u; i < arrayLength(&currentTrianglePos); i = i + 1) {
       if (i == index) {
         continue;
@@ -253,7 +253,7 @@ const mainCompute = tgpu['~unstable']
       + (alignment * params.alignmentStrength)
       + (cohesion * params.cohesionStrength);
     instanceInfo.velocity = normalize(instanceInfo.velocity) * clamp(length(instanceInfo.velocity), 0.0, 0.01);
-    
+
     if (instanceInfo.position[0] > 1.0 + triangleSize) {
       instanceInfo.position[0] = -1.0 - triangleSize;
     }
@@ -340,7 +340,7 @@ export const controls = {
     onButtonClick: () => paramsBuffer.write(presets.blobs),
   },
 
-  '⚛️ Particles': {
+  '⚛ Particles': {
     onButtonClick: () => paramsBuffer.write(presets.particles),
   },
 

apps/typegpu-docs/src/content/examples/simulation/confetti/index.ts

@@ -113,30 +113,24 @@ const mainVert = tgpu['~unstable']
     VertexOutput,
   )
   .does(
-    /* wgsl */ `(
-      @location(0) tilt: f32,
-      @location(1) angle: f32,
-      @location(2) color: vec4f,
-      @location(3) center: vec2f,
-      @builtin(vertex_index) index: u32,
-    ) -> VertexOutput {
-    let width = tilt;
-    let height = tilt / 2;
+    /* wgsl */ `(input: VertexInput) -> VertexOutput {
+    let width = input.tilt;
+    let height = input.tilt / 2;
 
     var pos = rotate(array<vec2f, 4>(
       vec2f(0, 0),
       vec2f(width, 0),
       vec2f(0, height),
       vec2f(width, height),
-    )[index] / 350, angle) + center;
+    )[input.index] / 350, input.angle) + input.center;
 
     if (canvasAspectRatio < 1) {
       pos.x /= canvasAspectRatio;
     } else {
       pos.y *= canvasAspectRatio;
     }
 
-    return VertexOutput(vec4f(pos, 0.0, 1.0), color);
+    return VertexOutput(vec4f(pos, 0.0, 1.0), input.color);
   }`,
   )
   .$uses({

apps/typegpu-docs/src/content/examples/simulation/fluid-double-buffering/index.ts

@@ -490,7 +490,7 @@ const vertexMain = tgpu['~unstable']
     { idx: d.builtin.vertexIndex },
     { pos: d.builtin.position, uv: d.vec2f },
   )
-  .does(/* wgsl */ `(@builtin(vertex_index) idx: u32) -> VertexOut {
+  .does(/* wgsl */ `(input: VertexInput) -> VertexOut {
     var pos = array<vec2f, 4>(
       vec2(1, 1), // top-right
       vec2(-1, 1), // top-left
@@ -506,8 +506,8 @@ const vertexMain = tgpu['~unstable']
     );
 
     var output: VertexOut;
-    output.pos = vec4f(pos[idx].x, pos[idx].y, 0.0, 1.0);
-    output.uv = uv[idx];
+    output.pos = vec4f(pos[input.idx].x, pos[input.idx].y, 0.0, 1.0);
+    output.uv = uv[input.idx];
     return output;
   }`);
 

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

@@ -58,3 +58,7 @@ export function createOutputType<T extends IOData>(returnType: IOLayout<T>) {
       : struct(withLocations(returnType) as Record<string, T>)
   ) as IOLayoutToOutputSchema<IOLayout<T>>;
 }
+
+export function createStructFromIO<T extends IOData>(members: IORecord<T>) {
+  return struct(withLocations(members) as Record<string, T>);
+}

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

@@ -1,6 +1,6 @@
 import type { OmitBuiltins } from '../../builtin';
-import { type Vec4f, isWgslStruct } from '../../data/wgslTypes';
-import type { TgpuNamable } from '../../namable';
+import type { Vec4f } from '../../data/wgslTypes';
+import { type TgpuNamable, isNamable } from '../../namable';
 import type { Labelled, ResolutionCtx, SelfResolvable } from '../../types';
 import { addReturnTypeToExternals } from '../resolve/externals';
 import { createFnCore } from './fnCore';
@@ -120,7 +120,7 @@ function createFragmentFn(
 
     $name(newLabel: string): This {
       core.label = newLabel;
-      if (isWgslStruct(outputType)) {
+      if (isNamable(outputType)) {
         outputType.$name(`${newLabel}_Output`);
       }
       return this;

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

@@ -1,6 +1,7 @@
 import type { OmitBuiltins } from '../../builtin';
-import { isWgslStruct } from '../../data/wgslTypes';
-import type { TgpuNamable } from '../../namable';
+import type { AnyWgslStruct } from '../../data/wgslTypes';
+import { type TgpuNamable, isNamable } from '../../namable';
+import type { GenerationCtx } from '../../smol/wgslGenerator';
 import type { Labelled, ResolutionCtx, SelfResolvable } from '../../types';
 import { addReturnTypeToExternals } from '../resolve/externals';
 import { createFnCore } from './fnCore';
@@ -11,7 +12,11 @@ import type {
   Implementation,
   InferIO,
 } from './fnTypes';
-import { type IOLayoutToOutputSchema, createOutputType } from './ioOutputType';
+import {
+  type IOLayoutToOutputSchema,
+  createOutputType,
+  createStructFromIO,
+} from './ioOutputType';
 
 // ----------
 // Public API
@@ -24,8 +29,9 @@ export interface TgpuVertexFnShell<
   VertexIn extends IOLayout,
   VertexOut extends IOLayout,
 > {
-  readonly argTypes: [VertexIn];
+  readonly argTypes: [AnyWgslStruct];
   readonly returnType: VertexOut;
+  readonly attributes: [VertexIn];
 
   /**
    * Creates a type-safe implementation of this signature
@@ -67,7 +73,7 @@ export interface TgpuVertexFn<
  *   passed onto the fragment shader stage.
  */
 export function vertexFn<
-  VertexIn extends IOLayout,
+  VertexIn extends IORecord,
   // Not allowing single-value output, as it is better practice
   // to properly label what the vertex shader is outputting.
   VertexOut extends IORecord,
@@ -76,8 +82,9 @@ export function vertexFn<
   outputType: VertexOut,
 ): TgpuVertexFnShell<ExoticIO<VertexIn>, ExoticIO<VertexOut>> {
   return {
-    argTypes: [inputType as ExoticIO<VertexIn>],
-    returnType: outputType as ExoticIO<VertexOut>,
+    attributes: [inputType as ExoticIO<VertexIn>],
+    returnType: createOutputType(outputType) as ExoticIO<VertexOut>,
+    argTypes: [createStructFromIO(inputType).$name('VertexInput')],
 
     does(implementation) {
       // biome-ignore lint/suspicious/noExplicitAny: <no thanks>
@@ -97,9 +104,9 @@ function createVertexFn(
   type This = TgpuVertexFn<IOLayout, IOLayout> & Labelled & SelfResolvable;
 
   const core = createFnCore(shell, implementation);
-  const outputType = createOutputType(shell.returnType);
+  const outputType = shell.returnType;
   if (typeof implementation === 'string') {
-    addReturnTypeToExternals(implementation, outputType, (externals) =>
+    addReturnTypeToExternals(implementation, shell.returnType, (externals) =>
       core.applyExternals(externals),
     );
   }
@@ -119,14 +126,29 @@ function createVertexFn(
 
     $name(newLabel: string): This {
       core.label = newLabel;
-      if (isWgslStruct(outputType)) {
+      if (isNamable(outputType)) {
         outputType.$name(`${newLabel}_Output`);
       }
       return this;
     },
 
     '~resolve'(ctx: ResolutionCtx): string {
-      return core.resolve(ctx, '@vertex ');
+      if (typeof implementation === 'string') {
+        return core.resolve(ctx, '@vertex ');
+      }
+
+      const generationCtx = ctx as GenerationCtx;
+      if (generationCtx.callStack === undefined) {
+        throw new Error(
+          'Cannot resolve a TGSL function outside of a generation context',
+        );
+      }
+
+      generationCtx.callStack.push(outputType);
+      const resolved = core.resolve(ctx, '@vertex ');
+      generationCtx.callStack.pop();
+
+      return resolved;
     },
 
     toString() {

packages/typegpu/src/core/pipeline/renderPipeline.ts

@@ -383,7 +383,7 @@ class RenderPipelineCore {
 
   constructor(public readonly options: RenderPipelineCoreOptions) {
     const connectedAttribs = connectAttributesToShader(
-      options.vertexFn.shell.argTypes[0],
+      options.vertexFn.shell.attributes[0],
       options.vertexAttribs,
     );
 

packages/typegpu/src/smol/wgslGenerator.ts

@@ -264,6 +264,15 @@ function generateStatement(
   }
 
   if ('r' in statement) {
+    // check if the thing at the top of the call stack is a struct
+    // if so wrap the value returned in a constructor of the struct (its resolved name)
+    if (
+      isWgslStruct(ctx.callStack[ctx.callStack.length - 1]) &&
+      statement.r !== null
+    ) {
+      return `${ctx.pre}return ${ctx.resolve(ctx.callStack[ctx.callStack.length - 1])}(${resolveRes(ctx, generateExpression(ctx, statement.r))});`;
+    }
+
     return statement.r === null
       ? `${ctx.pre}return;`
       : `${ctx.pre}return ${resolveRes(ctx, generateExpression(ctx, statement.r))};`;

packages/typegpu/tests/rawFn.test.ts

@@ -143,7 +143,7 @@ describe('tgpu.fn with raw string WGSL implementation', () => {
       template: `
         fn vs() {
           out.highlighted = highlighted.index;
-          
+
           let h = highlighted;
           let x = a.b.c.highlighted.d;
         }
@@ -285,7 +285,7 @@ struct fragment_Output {
     const func = tgpu['~unstable']
       .fn([d.vec4f, Point], d.vec2f)
       .does(/* wgsl */ `(
-        a: vec4f, 
+        a: vec4f,
         b : PointStruct ,
     ) -> vec2f {
     var newPoint: PointStruct;

packages/typegpu/tests/renderPipeline.test.ts

@@ -156,7 +156,7 @@ describe('Inter-Stage Variables', () => {
       .does('() { layout.bound.alpha; }')
       .$uses({ layout });
 
-    const fragmentFn = utgpu.vertexFn({}, { out: d.vec4f }).does('() {}');
+    const fragmentFn = utgpu.fragmentFn({}, { out: d.vec4f }).does('() {}');
 
     const pipeline = root
       .withVertex(vertexFn, {})

packages/typegpu/tests/tgslFn.test.ts

@@ -1,7 +1,8 @@
 import { parse } from 'tgpu-wgsl-parser';
 import { describe, expect, it } from 'vitest';
 import tgpu from '../src';
-import { f32, struct, vec3f } from '../src/data';
+import { builtin } from '../src/builtin';
+import { f32, struct, vec2f, vec3f, vec4f } from '../src/data';
 import { parseResolved } from './utils/parseResolved';
 
 describe('TGSL tgpu.fn function', () => {
@@ -162,4 +163,53 @@ describe('TGSL tgpu.fn function', () => {
 
     expect(actual).toEqual(expected);
   });
+
+  it('resolves vertexFn', () => {
+    const vertexFn = tgpu['~unstable']
+      .vertexFn(
+        {
+          vi: builtin.vertexIndex,
+          ii: builtin.instanceIndex,
+          color: vec4f,
+        },
+        {
+          pos: builtin.position,
+          uv: vec2f,
+        },
+      )
+      .does((input) => {
+        const vi = input.vi;
+        const ii = input.ii;
+        const color = input.color;
+
+        return {
+          pos: vec4f(color.w, ii, vi, 1),
+          uv: vec2f(color.w, vi),
+        };
+      })
+      .$name('vertex_fn');
+
+    const actual = parseResolved({ vertexFn });
+
+    const expected = parse(`
+      struct vertex_fn_Output {
+        @builtin(position) pos: vec4f,
+        @location(0) uv: vec2f,
+      }
+      struct VertexInput {
+        @builtin(vertex_index) vi: u32,
+        @builtin(instance_index) ii: u32,
+        @location(0) color: vec4f,
+      }
+
+      @vertex fn vertex_fn(input: VertexInput) -> vertex_fn_Output{
+        var vi = input.vi;
+        var ii = input.ii;
+        var color = input.color;
+        return vertex_fn_Output(vec4f(color.w, ii, vi, 1), vec2f(color.w, vi));
+      }
+    `);
+
+    expect(actual).toEqual(expected);
+  });
 });