From f30eb4b623855c63b166bf1805c5a1cade8c8265 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Fri, 17 Jan 2025 12:33:57 +0100
Subject: [PATCH 01/13] work work

---
 packages/tinyest-for-wgsl/src/parsers.ts   | 32 ++++++++++++-
 packages/tinyest/src/nodes.ts              | 12 ++++-
 packages/typegpu/src/data/struct.ts        | 12 ++++-
 packages/typegpu/src/resolutionCtx.ts      |  1 +
 packages/typegpu/src/smol/wgslGenerator.ts | 52 +++++++++++++++++++++-
 packages/typegpu/tests/struct.test.ts      | 15 ++++++-
 packages/typegpu/tests/tgslFn.test.ts      | 34 +++++++++++++-
 7 files changed, 151 insertions(+), 7 deletions(-)

diff --git a/packages/tinyest-for-wgsl/src/parsers.ts b/packages/tinyest-for-wgsl/src/parsers.ts
index 0d44a2f4c..95f72256e 100644
--- a/packages/tinyest-for-wgsl/src/parsers.ts
+++ b/packages/tinyest-for-wgsl/src/parsers.ts
@@ -185,7 +185,7 @@ const Transpilers: Partial<{
 
   Literal(ctx, node) {
     if (typeof node.value === 'string') {
-      throw new Error('String literals are not supported in TGSL.');
+      return { s: node.value };
     }
     return { n: node.raw ?? '' };
   },
@@ -249,6 +249,36 @@ const Transpilers: Partial<{
       q: alternate ? [test, consequent, alternate] : [test, consequent],
     };
   },
+
+  ObjectExpression(ctx, node) {
+    const properties: Record<string, smol.Expression> = {};
+
+    for (const prop of node.properties) {
+      // TODO: Handle SpreadElement
+      if (prop.type === 'SpreadElement') {
+        throw new Error('Spread elements are not supported in TGSL.');
+      }
+
+      // TODO: Handle computed properties
+      if (prop.key.type !== 'Identifier' && prop.key.type !== 'Literal') {
+        throw new Error(
+          'Only Identifier and Literal keys are supported as object keys.',
+        );
+      }
+
+      ctx.ignoreExternalDepth++;
+      const key =
+        prop.key.type === 'Identifier'
+          ? (transpile(ctx, prop.key) as string)
+          : String(prop.key.value);
+      const value = transpile(ctx, prop.value) as smol.Expression;
+      ctx.ignoreExternalDepth--;
+
+      properties[key] = value;
+    }
+
+    return { o: properties };
+  },
 };
 
 function transpile(ctx: Context, node: acorn.AnyNode): smol.AnyNode {
diff --git a/packages/tinyest/src/nodes.ts b/packages/tinyest/src/nodes.ts
index b4ec754be..194890481 100644
--- a/packages/tinyest/src/nodes.ts
+++ b/packages/tinyest/src/nodes.ts
@@ -109,6 +109,10 @@ export type UnaryExpression = {
   u: [op: UnaryOperator, inner: Expression];
 };
 
+export type ObjectExpression = {
+  o: Record<string, Expression>;
+};
+
 export type MemberAccess = {
   a: [object: Expression, member: string];
 };
@@ -126,7 +130,12 @@ export type Num = {
   n: string;
 };
 
-export type Literal = Num | boolean;
+/** A string literal */
+export type Str = {
+  s: string;
+};
+
+export type Literal = Num | Str | boolean;
 
 /** Identifiers are just strings, since string literals are rare in WGSL, and identifiers are everywhere. */
 export type Expression =
@@ -135,6 +144,7 @@ export type Expression =
   | AssignmentExpression
   | LogicalExpression
   | UnaryExpression
+  | ObjectExpression
   | MemberAccess
   | IndexAccess
   | Call
diff --git a/packages/typegpu/src/data/struct.ts b/packages/typegpu/src/data/struct.ts
index 329a4edbf..c30d214b5 100644
--- a/packages/typegpu/src/data/struct.ts
+++ b/packages/typegpu/src/data/struct.ts
@@ -21,6 +21,10 @@ export interface TgpuStruct<TProps extends Record<string, BaseWgslData>>
   readonly '~exotic': WgslStruct<ExoticRecord<TProps>>;
 }
 
+export type NativeStruct<TProps extends Record<string, AnyWgslData>> =
+  TgpuStruct<Prettify<TProps>> &
+    ((props: InferRecord<TProps>) => InferRecord<TProps>);
+
 /**
  * Creates a struct schema that can be used to construct GPU buffers.
  * Ensures proper alignment and padding of properties (as opposed to a `d.unstruct` schema).
@@ -34,8 +38,12 @@ export interface TgpuStruct<TProps extends Record<string, BaseWgslData>>
  */
 export const struct = <TProps extends Record<string, AnyWgslData>>(
   props: TProps,
-): TgpuStruct<Prettify<ExoticRecord<TProps>>> =>
-  new TgpuStructImpl(props as ExoticRecord<TProps>);
+): NativeStruct<ExoticRecord<TProps>> => {
+  const struct = new TgpuStructImpl(props as ExoticRecord<TProps>);
+  const construct = (props: InferRecord<ExoticRecord<TProps>>) => props;
+
+  return Object.assign(construct, struct);
+};
 
 // --------------
 // Implementation
diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts
index a96440382..f436bf753 100644
--- a/packages/typegpu/src/resolutionCtx.ts
+++ b/packages/typegpu/src/resolutionCtx.ts
@@ -264,6 +264,7 @@ class ResolutionCtxImpl implements ResolutionCtx {
   public readonly fixedBindings: FixedBindingConfig[] = [];
   // --
 
+  public readonly callStack: unknown[] = [];
   public readonly names: NameRegistry;
 
   constructor(opts: ResolutionCtxImplOptions) {
diff --git a/packages/typegpu/src/smol/wgslGenerator.ts b/packages/typegpu/src/smol/wgslGenerator.ts
index 166a7772c..6c311a403 100644
--- a/packages/typegpu/src/smol/wgslGenerator.ts
+++ b/packages/typegpu/src/smol/wgslGenerator.ts
@@ -1,6 +1,6 @@
 import type * as smol from 'tinyest';
 import { bool } from '../data';
-import { isWgslData } from '../data/wgslTypes';
+import { isWgslData, isWgslStruct } from '../data/wgslTypes';
 import {
   type ResolutionCtx,
   type Resource,
@@ -32,6 +32,7 @@ const parenthesizedOps = [
 
 export type GenerationCtx = ResolutionCtx & {
   readonly pre: string;
+  readonly callStack: unknown[];
   indent(): string;
   dedent(): string;
   getById(id: string): Resource;
@@ -179,9 +180,13 @@ function generateExpression(
     const id = generateExpression(ctx, callee);
     const idValue = id.value;
 
+    ctx.callStack.push(idValue);
+
     const argResources = args.map((arg) => generateExpression(ctx, arg));
     const argValues = argResources.map((res) => resolveRes(ctx, res));
 
+    ctx.callStack.pop();
+
     if (typeof idValue === 'string') {
       return {
         value: `${idValue}(${argValues.join(', ')})`,
@@ -189,6 +194,15 @@ function generateExpression(
       };
     }
 
+    if (isWgslStruct(idValue)) {
+      const id = ctx.resolve(idValue);
+
+      return {
+        value: `${id}(${argValues.join(', ')})`,
+        dataType: UnknownData,
+      };
+    }
+
     // Assuming that `id` is callable
     // TODO: Pass in resources, not just values.
     const result = (idValue as unknown as (...args: unknown[]) => unknown)(
@@ -198,6 +212,42 @@ function generateExpression(
     return { value: result, dataType: UnknownData };
   }
 
+  if ('o' in expression) {
+    const obj = expression.o;
+    const callee = ctx.callStack[ctx.callStack.length - 1];
+
+    const generateEntries = (values: smol.Expression[]) =>
+      values
+        .map((value) => {
+          const valueRes = generateExpression(ctx, value);
+          return resolveRes(ctx, valueRes);
+        })
+        .join(', ');
+
+    if (isWgslStruct(callee)) {
+      const propKeys = Object.keys(callee.propTypes);
+      const values = propKeys.map((key) => {
+        const val = obj[key];
+        if (val === undefined) {
+          throw new Error(
+            `Missing property ${key} in object literal for struct ${callee}`,
+          );
+        }
+        return val;
+      });
+
+      return {
+        value: generateEntries(values),
+        dataType: UnknownData,
+      };
+    }
+
+    return {
+      value: generateEntries(Object.values(obj)),
+      dataType: UnknownData,
+    };
+  }
+
   assertExhaustive(expression);
 }
 
diff --git a/packages/typegpu/tests/struct.test.ts b/packages/typegpu/tests/struct.test.ts
index e3d844091..eede8f2a2 100644
--- a/packages/typegpu/tests/struct.test.ts
+++ b/packages/typegpu/tests/struct.test.ts
@@ -1,5 +1,5 @@
 import { BufferReader, BufferWriter } from 'typed-binary';
-import { describe, expect, it } from 'vitest';
+import { describe, expect, expectTypeOf, it } from 'vitest';
 import {
   alignmentOf,
   arrayOf,
@@ -9,6 +9,7 @@ import {
   sizeOf,
   struct,
   u32,
+  type v3u,
   vec2h,
   vec2u,
   vec3f,
@@ -224,4 +225,16 @@ describe('struct', () => {
     writeData(new BufferWriter(buffer2), TestStruct2, value2);
     expect(readData(new BufferReader(buffer2), TestStruct2)).toEqual(value2);
   });
+
+  it('can be called to create an object', () => {
+    const TestStruct = struct({
+      x: u32,
+      y: vec3u,
+    });
+
+    const obj = TestStruct({ x: 1, y: vec3u(1, 2, 3) });
+
+    expect(obj).toEqual({ x: 1, y: vec3u(1, 2, 3) });
+    expectTypeOf(obj).toEqualTypeOf<{ x: number; y: v3u }>();
+  });
 });
diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts
index 2fe75ce3d..314f5ef28 100644
--- a/packages/typegpu/tests/tgslFn.test.ts
+++ b/packages/typegpu/tests/tgslFn.test.ts
@@ -1,6 +1,6 @@
 import { parse } from 'tgpu-wgsl-parser';
 import { describe, expect, it } from 'vitest';
-import { f32, vec3f } from '../src/data';
+import { f32, struct, vec3f } from '../src/data';
 import tgpu from '../src/experimental';
 import { parseResolved } from './utils/parseResolved';
 
@@ -86,4 +86,36 @@ describe('TGSL tgpu.fn function', () => {
 
     expect(actual).toEqual(expected);
   });
+
+  it('resolves structs', () => {
+    const Gradient = struct({
+      from: vec3f,
+      to: vec3f,
+    });
+
+    const createGradient = tgpu
+      .fn([], Gradient)
+      .does(() => {
+        return Gradient({ to: vec3f(1, 2, 3), from: vec3f(4, 5, 6) });
+      })
+      .$name('create_gradient');
+
+    const actual = parseResolved({ createGradient });
+
+    const expected = parse(`
+      struct Gradient {
+        from: vec3f,
+        to: vec3f,
+      }
+
+      fn create_gradient() -> Gradient {
+        return Gradient(vec3f(4, 5, 6), vec3f(1, 2, 3));
+      }
+    `);
+
+    console.log(actual);
+    console.log(expected);
+
+    expect(actual).toEqual(expected);
+  });
 });

From 64216ce205007c6b34a42621e919c8f835eb02e1 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Fri, 17 Jan 2025 16:35:18 +0100
Subject: [PATCH 02/13] new approach

---
 packages/typegpu/src/data/sizeOf.ts    | 12 +++++--
 packages/typegpu/src/data/struct.ts    | 47 +++++++++++---------------
 packages/typegpu/src/data/wgslTypes.ts |  1 +
 packages/typegpu/tests/tgslFn.test.ts  |  3 --
 4 files changed, 31 insertions(+), 32 deletions(-)

diff --git a/packages/typegpu/src/data/sizeOf.ts b/packages/typegpu/src/data/sizeOf.ts
index 167ed4f4f..f5bc14cff 100644
--- a/packages/typegpu/src/data/sizeOf.ts
+++ b/packages/typegpu/src/data/sizeOf.ts
@@ -1,3 +1,4 @@
+import type { TgpuStruct } from '.';
 import { roundUp } from '../mathUtils';
 import { alignmentOf, customAlignmentOf } from './alignmentOf';
 import type { AnyData, LooseTypeLiteral, Unstruct } from './dataTypes';
@@ -7,7 +8,12 @@ import {
   isLooseDecorated,
   isUnstruct,
 } from './dataTypes';
-import type { BaseWgslData, WgslStruct, WgslTypeLiteral } from './wgslTypes';
+import type {
+  AnyWgslData,
+  BaseWgslData,
+  WgslStruct,
+  WgslTypeLiteral,
+} from './wgslTypes';
 import { isDecorated, isWgslArray, isWgslStruct } from './wgslTypes';
 
 const knownSizesMap: Record<string, number> = {
@@ -163,6 +169,8 @@ export function sizeOf(schema: BaseWgslData): number {
 /**
  * Returns the size (in bytes) of data represented by the `schema`.
  */
-export function PUBLIC_sizeOf(schema: AnyData): number {
+export function PUBLIC_sizeOf<T extends Record<string, AnyWgslData>>(
+  schema: AnyData | TgpuStruct<T>,
+): number {
   return sizeOf(schema);
 }
diff --git a/packages/typegpu/src/data/struct.ts b/packages/typegpu/src/data/struct.ts
index c30d214b5..90fa020b9 100644
--- a/packages/typegpu/src/data/struct.ts
+++ b/packages/typegpu/src/data/struct.ts
@@ -18,13 +18,9 @@ import type { AnyWgslData, BaseWgslData, WgslStruct } from './wgslTypes';
 export interface TgpuStruct<TProps extends Record<string, BaseWgslData>>
   extends WgslStruct<TProps>,
     TgpuNamable {
-  readonly '~exotic': WgslStruct<ExoticRecord<TProps>>;
+  readonly '~exotic': WgslStruct<TProps>;
 }
 
-export type NativeStruct<TProps extends Record<string, AnyWgslData>> =
-  TgpuStruct<Prettify<TProps>> &
-    ((props: InferRecord<TProps>) => InferRecord<TProps>);
-
 /**
  * Creates a struct schema that can be used to construct GPU buffers.
  * Ensures proper alignment and padding of properties (as opposed to a `d.unstruct` schema).
@@ -38,40 +34,37 @@ export type NativeStruct<TProps extends Record<string, AnyWgslData>> =
  */
 export const struct = <TProps extends Record<string, AnyWgslData>>(
   props: TProps,
-): NativeStruct<ExoticRecord<TProps>> => {
-  const struct = new TgpuStructImpl(props as ExoticRecord<TProps>);
-  const construct = (props: InferRecord<ExoticRecord<TProps>>) => props;
+): TgpuStruct<Prettify<ExoticRecord<TProps>>> => {
+  const struct = <T>(props: T) => {
+    return props;
+  };
+  Object.setPrototypeOf(struct, TgpuStructImpl);
+  struct.propTypes = props as ExoticRecord<TProps>;
+  struct['~repr'] = {} as InferRecord<TProps>;
+  struct['~exotic'] = {} as WgslStruct<ExoticRecord<TProps>>;
 
-  return Object.assign(construct, struct);
+  return struct as unknown as TgpuStruct<Prettify<ExoticRecord<TProps>>>;
 };
 
 // --------------
 // Implementation
 // --------------
 
-class TgpuStructImpl<TProps extends Record<string, AnyWgslData>>
-  implements TgpuStruct<TProps>
-{
-  private _label: string | undefined;
-
-  public readonly type = 'struct';
-  /** Type-token, not available at runtime */
-  public readonly '~repr'!: InferRecord<TProps>;
-  /** Type-token, not available at runtime */
-  public readonly '~exotic'!: WgslStruct<ExoticRecord<TProps>>;
-
-  constructor(public readonly propTypes: TProps) {}
+const TgpuStructImpl = {
+  type: 'struct',
 
-  get label() {
+  get label(): string | undefined {
+    //@ts-ignore
     return this._label;
-  }
+  },
 
   $name(label: string) {
+    //@ts-ignore
     this._label = label;
     return this;
-  }
+  },
 
-  toString() {
+  toString(): string {
     return `struct:${this.label ?? '<unnamed>'}`;
-  }
-}
+  },
+};
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index cd685d58e..1aa10b03d 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -675,6 +675,7 @@ export interface Mat4x4f {
 export interface WgslStruct<
   TProps extends Record<string, BaseWgslData> = Record<string, BaseWgslData>,
 > {
+  <T>(props: T): T;
   readonly type: 'struct';
   readonly label?: string | undefined;
   readonly propTypes: TProps;
diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts
index 314f5ef28..0b3bd0a58 100644
--- a/packages/typegpu/tests/tgslFn.test.ts
+++ b/packages/typegpu/tests/tgslFn.test.ts
@@ -113,9 +113,6 @@ describe('TGSL tgpu.fn function', () => {
       }
     `);
 
-    console.log(actual);
-    console.log(expected);
-
     expect(actual).toEqual(expected);
   });
 });

From 752adea646861f6dc37cbf9d5a002fbc0046003d Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Fri, 17 Jan 2025 17:43:01 +0100
Subject: [PATCH 03/13] i'm sorry TypeScript

---
 packages/typegpu/src/data/sizeOf.ts    | 12 ++----------
 packages/typegpu/src/data/wgslTypes.ts |  5 ++++-
 2 files changed, 6 insertions(+), 11 deletions(-)

diff --git a/packages/typegpu/src/data/sizeOf.ts b/packages/typegpu/src/data/sizeOf.ts
index f5bc14cff..167ed4f4f 100644
--- a/packages/typegpu/src/data/sizeOf.ts
+++ b/packages/typegpu/src/data/sizeOf.ts
@@ -1,4 +1,3 @@
-import type { TgpuStruct } from '.';
 import { roundUp } from '../mathUtils';
 import { alignmentOf, customAlignmentOf } from './alignmentOf';
 import type { AnyData, LooseTypeLiteral, Unstruct } from './dataTypes';
@@ -8,12 +7,7 @@ import {
   isLooseDecorated,
   isUnstruct,
 } from './dataTypes';
-import type {
-  AnyWgslData,
-  BaseWgslData,
-  WgslStruct,
-  WgslTypeLiteral,
-} from './wgslTypes';
+import type { BaseWgslData, WgslStruct, WgslTypeLiteral } from './wgslTypes';
 import { isDecorated, isWgslArray, isWgslStruct } from './wgslTypes';
 
 const knownSizesMap: Record<string, number> = {
@@ -169,8 +163,6 @@ export function sizeOf(schema: BaseWgslData): number {
 /**
  * Returns the size (in bytes) of data represented by the `schema`.
  */
-export function PUBLIC_sizeOf<T extends Record<string, AnyWgslData>>(
-  schema: AnyData | TgpuStruct<T>,
-): number {
+export function PUBLIC_sizeOf(schema: AnyData): number {
   return sizeOf(schema);
 }
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index 1aa10b03d..bb0601b94 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -683,6 +683,9 @@ export interface WgslStruct<
   readonly '~repr': InferRecord<TProps>;
 }
 
+// biome-ignore lint/suspicious/noExplicitAny: <i'm sorry>
+export type AnyWgslStruct = WgslStruct<any>;
+
 export interface WgslArray<TElement = BaseWgslData> {
   readonly type: 'array';
   readonly elementCount: number;
@@ -815,7 +818,7 @@ export type AnyWgslData =
   | Mat2x2f
   | Mat3x3f
   | Mat4x4f
-  | WgslStruct
+  | AnyWgslStruct
   | WgslArray
   | Atomic
   | Decorated;

From c054d3c62a116a8968c0e58536fa971115a27ed8 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Fri, 17 Jan 2025 17:58:37 +0100
Subject: [PATCH 04/13] nested test

---
 packages/typegpu/tests/tgslFn.test.ts | 47 +++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts
index 8ab98ab99..0e4d1883b 100644
--- a/packages/typegpu/tests/tgslFn.test.ts
+++ b/packages/typegpu/tests/tgslFn.test.ts
@@ -115,4 +115,51 @@ describe('TGSL tgpu.fn function', () => {
 
     expect(actual).toEqual(expected);
   });
+
+  it('resolves deeply nested structs', () => {
+    const A = struct({
+      b: f32,
+    }).$name('A');
+
+    const B = struct({
+      a: A,
+      c: f32,
+    }).$name('B');
+
+    const C = struct({
+      b: B,
+      a: A,
+    }).$name('C');
+
+    const pureConfusion = tgpu['~unstable']
+      .fn([], A)
+      .does(() => {
+        return C({ a: A({ b: 3 }), b: B({ a: A({ b: 4 }), c: 5 }) }).a;
+      })
+      .$name('pure_confusion');
+
+    const actual = parseResolved({ pureConfusion });
+
+    const expected = parse(`
+      struct A {
+        b: f32,
+      }
+
+      struct B {
+        a: A,
+        c: f32,
+      }
+
+      struct C {
+        b: B,
+        a: A,
+      }
+
+      fn pure_confusion() -> A {
+        return C(B(A(4), 5), A(3)).a;
+      }
+    `);
+
+    expect(actual).toEqual(expected);
+  });
 });

From bf2d43e37fed0d3430690262e726fe25358e763d Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 08:34:30 +0100
Subject: [PATCH 05/13] better label handling

---
 packages/typegpu/src/data/struct.ts | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/packages/typegpu/src/data/struct.ts b/packages/typegpu/src/data/struct.ts
index 90fa020b9..6f278b418 100644
--- a/packages/typegpu/src/data/struct.ts
+++ b/packages/typegpu/src/data/struct.ts
@@ -52,14 +52,13 @@ export const struct = <TProps extends Record<string, AnyWgslData>>(
 
 const TgpuStructImpl = {
   type: 'struct',
+  _label: undefined as string | undefined,
 
   get label(): string | undefined {
-    //@ts-ignore
     return this._label;
   },
 
   $name(label: string) {
-    //@ts-ignore
     this._label = label;
     return this;
   },

From 711813c09ee743df85d0618f25d9ca8f62833041 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 08:48:53 +0100
Subject: [PATCH 06/13] fix vertex attributes

---
 packages/typegpu/src/core/vertexLayout/vertexAttribute.ts | 4 ++--
 packages/typegpu/src/data/wgslTypes.ts                    | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/packages/typegpu/src/core/vertexLayout/vertexAttribute.ts b/packages/typegpu/src/core/vertexLayout/vertexAttribute.ts
index 2ad54ceaf..a05f4f12d 100644
--- a/packages/typegpu/src/core/vertexLayout/vertexAttribute.ts
+++ b/packages/typegpu/src/core/vertexLayout/vertexAttribute.ts
@@ -1,5 +1,5 @@
 import type { Disarray, Unstruct } from '../../data/dataTypes';
-import type { WgslArray, WgslStruct } from '../../data/wgslTypes';
+import type { AnyWgslStruct, WgslArray } from '../../data/wgslTypes';
 import type {
   KindToAcceptedAttribMap,
   KindToDefaultFormatMap,
@@ -15,7 +15,7 @@ import type {
  * - TgpuStruct<{ a: Vec3f, b: unorm8x2 }>
  * - TgpuStruct<{ nested: TgpuStruct<{ a: Vec3f }> }>
  */
-export type DataToContainedAttribs<T> = T extends WgslStruct | Unstruct
+export type DataToContainedAttribs<T> = T extends AnyWgslStruct | Unstruct
   ? {
       [Key in keyof T['propTypes']]: DataToContainedAttribs<
         T['propTypes'][Key]
diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index bb0601b94..369e27b57 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -675,7 +675,7 @@ export interface Mat4x4f {
 export interface WgslStruct<
   TProps extends Record<string, BaseWgslData> = Record<string, BaseWgslData>,
 > {
-  <T>(props: T): T;
+  (props: InferRecord<TProps>): InferRecord<TProps>;
   readonly type: 'struct';
   readonly label?: string | undefined;
   readonly propTypes: TProps;

From 3b397052989c781f6d9876fb1d3385def70f4410 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 08:51:32 +0100
Subject: [PATCH 07/13] better explanation

---
 packages/typegpu/src/data/wgslTypes.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts
index 369e27b57..865d17e93 100644
--- a/packages/typegpu/src/data/wgslTypes.ts
+++ b/packages/typegpu/src/data/wgslTypes.ts
@@ -683,7 +683,7 @@ export interface WgslStruct<
   readonly '~repr': InferRecord<TProps>;
 }
 
-// biome-ignore lint/suspicious/noExplicitAny: <i'm sorry>
+// biome-ignore lint/suspicious/noExplicitAny: <we need the type to be broader than WgslStruct<Record<string, BaseWgslData>>
 export type AnyWgslStruct = WgslStruct<any>;
 
 export interface WgslArray<TElement = BaseWgslData> {

From bf12678f4d84fb899863a7396833470a11a565e9 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 20 Jan 2025 08:53:59 +0100
Subject: [PATCH 08/13] additional test

---
 packages/typegpu/tests/struct.test.ts | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/packages/typegpu/tests/struct.test.ts b/packages/typegpu/tests/struct.test.ts
index eede8f2a2..be3ec0e31 100644
--- a/packages/typegpu/tests/struct.test.ts
+++ b/packages/typegpu/tests/struct.test.ts
@@ -237,4 +237,14 @@ describe('struct', () => {
     expect(obj).toEqual({ x: 1, y: vec3u(1, 2, 3) });
     expectTypeOf(obj).toEqualTypeOf<{ x: number; y: v3u }>();
   });
+
+  it('cannot be called with invalid properties', () => {
+    const TestStruct = struct({
+      x: u32,
+      y: vec3u,
+    });
+
+    // @ts-expect-error
+    TestStruct({ x: 1, z: 2 });
+  });
 });

From 29afffa47c3e0ba60befe85392e1b55ba2048baf Mon Sep 17 00:00:00 2001
From: reczkok <66403540+reczkok@users.noreply.github.com>
Date: Mon, 20 Jan 2025 11:15:47 +0100
Subject: [PATCH 09/13] Yes

Co-authored-by: Marcin Hawryluk <70582973+mhawryluk@users.noreply.github.com>
---
 packages/typegpu/src/data/struct.ts | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/packages/typegpu/src/data/struct.ts b/packages/typegpu/src/data/struct.ts
index 6f278b418..8ffe09c3d 100644
--- a/packages/typegpu/src/data/struct.ts
+++ b/packages/typegpu/src/data/struct.ts
@@ -35,9 +35,7 @@ export interface TgpuStruct<TProps extends Record<string, BaseWgslData>>
 export const struct = <TProps extends Record<string, AnyWgslData>>(
   props: TProps,
 ): TgpuStruct<Prettify<ExoticRecord<TProps>>> => {
-  const struct = <T>(props: T) => {
-    return props;
-  };
+  const struct = <T>(props: T) => props;
   Object.setPrototypeOf(struct, TgpuStructImpl);
   struct.propTypes = props as ExoticRecord<TProps>;
   struct['~repr'] = {} as InferRecord<TProps>;

From cbe38d2be4b2bb0cf05e8ec58c8060b45893d8e5 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Wed, 22 Jan 2025 11:26:45 +0100
Subject: [PATCH 10/13] creating structs

---
 .../typegpu/src/core/function/ioOutputType.ts |  4 ++
 .../src/core/function/tgpuFragmentFn.ts       |  6 +--
 .../typegpu/src/core/function/tgpuVertexFn.ts | 26 +++++++----
 .../src/core/pipeline/renderPipeline.ts       |  2 +-
 packages/typegpu/tests/rawFn.test.ts          |  4 +-
 packages/typegpu/tests/tgslFn.test.ts         | 46 ++++++++++++++++++-
 6 files changed, 71 insertions(+), 17 deletions(-)

diff --git a/packages/typegpu/src/core/function/ioOutputType.ts b/packages/typegpu/src/core/function/ioOutputType.ts
index 506b6d922..15ab0a82d 100644
--- a/packages/typegpu/src/core/function/ioOutputType.ts
+++ b/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>);
+}
diff --git a/packages/typegpu/src/core/function/tgpuFragmentFn.ts b/packages/typegpu/src/core/function/tgpuFragmentFn.ts
index f824928b7..0ee7739b2 100644
--- a/packages/typegpu/src/core/function/tgpuFragmentFn.ts
+++ b/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;
diff --git a/packages/typegpu/src/core/function/tgpuVertexFn.ts b/packages/typegpu/src/core/function/tgpuVertexFn.ts
index cb7512940..577a35123 100644
--- a/packages/typegpu/src/core/function/tgpuVertexFn.ts
+++ b/packages/typegpu/src/core/function/tgpuVertexFn.ts
@@ -1,6 +1,6 @@
 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 { Labelled, ResolutionCtx, SelfResolvable } from '../../types';
 import { addReturnTypeToExternals } from '../resolve/externals';
 import { createFnCore } from './fnCore';
@@ -11,7 +11,11 @@ import type {
   Implementation,
   InferIO,
 } from './fnTypes';
-import { type IOLayoutToOutputSchema, createOutputType } from './ioOutputType';
+import {
+  type IOLayoutToOutputSchema,
+  createOutputType,
+  createStructFromIO,
+} from './ioOutputType';
 
 // ----------
 // Public API
@@ -24,8 +28,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 +72,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 +81,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 +103,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,7 +125,7 @@ function createVertexFn(
 
     $name(newLabel: string): This {
       core.label = newLabel;
-      if (isWgslStruct(outputType)) {
+      if (isNamable(outputType)) {
         outputType.$name(`${newLabel}_Output`);
       }
       return this;
diff --git a/packages/typegpu/src/core/pipeline/renderPipeline.ts b/packages/typegpu/src/core/pipeline/renderPipeline.ts
index 30422c6bd..eb44c9163 100644
--- a/packages/typegpu/src/core/pipeline/renderPipeline.ts
+++ b/packages/typegpu/src/core/pipeline/renderPipeline.ts
@@ -280,7 +280,7 @@ class RenderPipelineCore {
 
   constructor(public readonly options: RenderPipelineCoreOptions) {
     const connectedAttribs = connectAttributesToShader(
-      options.vertexFn.shell.argTypes[0],
+      options.vertexFn.shell.attributes[0],
       options.vertexAttribs,
     );
 
diff --git a/packages/typegpu/tests/rawFn.test.ts b/packages/typegpu/tests/rawFn.test.ts
index c9b9b14df..dd394926f 100644
--- a/packages/typegpu/tests/rawFn.test.ts
+++ b/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;
diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts
index 0e4d1883b..1a2b9a6f7 100644
--- a/packages/typegpu/tests/tgslFn.test.ts
+++ b/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,47 @@ 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(`
+      fn vertex_fn(
+        vi: i32,
+        ii: i32,
+        color: vec4<f32>,
+      ) -> Output {
+        return Output(
+          vec4<f32>(color.w, ii, vi, 1),
+          vec2<f32>(color.w, vi),
+        );
+      }
+    `);
+
+    expect(actual).toEqual(expected);
+  });
 });

From bb9103174b63f39d344566567333742b8c86b27f Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Wed, 22 Jan 2025 12:17:47 +0100
Subject: [PATCH 11/13] maybe better call stack

---
 .../typegpu/src/core/function/tgpuVertexFn.ts | 18 ++++++++++++-
 packages/typegpu/src/smol/wgslGenerator.ts    |  9 +++++++
 packages/typegpu/tests/tgslFn.test.ts         | 26 ++++++++++++-------
 3 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/packages/typegpu/src/core/function/tgpuVertexFn.ts b/packages/typegpu/src/core/function/tgpuVertexFn.ts
index 577a35123..204266a20 100644
--- a/packages/typegpu/src/core/function/tgpuVertexFn.ts
+++ b/packages/typegpu/src/core/function/tgpuVertexFn.ts
@@ -1,6 +1,7 @@
 import type { OmitBuiltins } from '../../builtin';
 import type { AnyWgslStruct } from '../../data/wgslTypes';
 import { type TgpuNamable, isNamable } from '../../namable';
+import type { GenerationCtx } from '../../smol';
 import type { Labelled, ResolutionCtx, SelfResolvable } from '../../types';
 import { addReturnTypeToExternals } from '../resolve/externals';
 import { createFnCore } from './fnCore';
@@ -132,7 +133,22 @@ function createVertexFn(
     },
 
     '~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() {
diff --git a/packages/typegpu/src/smol/wgslGenerator.ts b/packages/typegpu/src/smol/wgslGenerator.ts
index 6c311a403..214fcdcb1 100644
--- a/packages/typegpu/src/smol/wgslGenerator.ts
+++ b/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))};`;
diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts
index 1a2b9a6f7..d5485d4cd 100644
--- a/packages/typegpu/tests/tgslFn.test.ts
+++ b/packages/typegpu/tests/tgslFn.test.ts
@@ -141,6 +141,8 @@ describe('TGSL tgpu.fn function', () => {
 
     const actual = parseResolved({ pureConfusion });
 
+    console.log(tgpu.resolve({ externals: { pureConfusion } }));
+
     const expected = parse(`
       struct A {
         b: f32,
@@ -192,15 +194,21 @@ describe('TGSL tgpu.fn function', () => {
     const actual = parseResolved({ vertexFn });
 
     const expected = parse(`
-      fn vertex_fn(
-        vi: i32,
-        ii: i32,
-        color: vec4<f32>,
-      ) -> Output {
-        return Output(
-          vec4<f32>(color.w, ii, vi, 1),
-          vec2<f32>(color.w, vi),
-        );
+      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));
       }
     `);
 

From e9391acdab6b2fc59452cb9258fea27aeda33522 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Mon, 27 Jan 2025 15:11:33 +0100
Subject: [PATCH 12/13] fix import

---
 packages/typegpu/src/core/function/tgpuVertexFn.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/typegpu/src/core/function/tgpuVertexFn.ts b/packages/typegpu/src/core/function/tgpuVertexFn.ts
index 204266a20..292117c58 100644
--- a/packages/typegpu/src/core/function/tgpuVertexFn.ts
+++ b/packages/typegpu/src/core/function/tgpuVertexFn.ts
@@ -1,7 +1,7 @@
 import type { OmitBuiltins } from '../../builtin';
 import type { AnyWgslStruct } from '../../data/wgslTypes';
 import { type TgpuNamable, isNamable } from '../../namable';
-import type { GenerationCtx } from '../../smol';
+import type { GenerationCtx } from '../../smol/wgslGenerator';
 import type { Labelled, ResolutionCtx, SelfResolvable } from '../../types';
 import { addReturnTypeToExternals } from '../resolve/externals';
 import { createFnCore } from './fnCore';

From 32af044e9e6770c62b35a0caca07f698eeace5f5 Mon Sep 17 00:00:00 2001
From: reczkok <reczkok@gmail.com>
Date: Wed, 29 Jan 2025 12:33:08 +0100
Subject: [PATCH 13/13] adapt examples

---
 .../examples/rendering/box-raytracing/index.ts   |  4 ++--
 .../examples/simulation/boids-next/index.ts      | 14 +++++++-------
 .../examples/simulation/confetti/index.ts        | 16 +++++-----------
 .../simulation/fluid-double-buffering/index.ts   |  6 +++---
 4 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/apps/typegpu-docs/src/content/examples/rendering/box-raytracing/index.ts b/apps/typegpu-docs/src/content/examples/rendering/box-raytracing/index.ts
index 6658a1db6..9c91f0997 100644
--- a/apps/typegpu-docs/src/content/examples/rendering/box-raytracing/index.ts
+++ b/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');
diff --git a/apps/typegpu-docs/src/content/examples/simulation/boids-next/index.ts b/apps/typegpu-docs/src/content/examples/simulation/boids-next/index.ts
index d43b67c25..a5a566560 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/boids-next/index.ts
+++ b/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),
   },
 
diff --git a/apps/typegpu-docs/src/content/examples/simulation/confetti/index.ts b/apps/typegpu-docs/src/content/examples/simulation/confetti/index.ts
index 5e70000b8..64e164fe5 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/confetti/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/confetti/index.ts
@@ -113,22 +113,16 @@ 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;
@@ -136,7 +130,7 @@ const mainVert = tgpu['~unstable']
       pos.y *= canvasAspectRatio;
     }
 
-    return VertexOutput(vec4f(pos, 0.0, 1.0), color);
+    return VertexOutput(vec4f(pos, 0.0, 1.0), input.color);
   }`,
   )
   .$uses({
diff --git a/apps/typegpu-docs/src/content/examples/simulation/fluid-double-buffering/index.ts b/apps/typegpu-docs/src/content/examples/simulation/fluid-double-buffering/index.ts
index 163353759..ebce98f44 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/fluid-double-buffering/index.ts
+++ b/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;
   }`);