feat: More convenient TgpuVertexFn.AutoIn and TgpuVertexFn.AutoOut types (#2282)

Author: iwoplaza

apps/typegpu-docs/src/examples/algorithms/jump-flood-voronoi/index.ts

@@ -6,7 +6,6 @@ import { defineControls } from '../../common/defineControls.ts';
 const root = await tgpu.init();
 
 const canvas = document.querySelector('canvas') as HTMLCanvasElement;
-const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
 
 const context = root.configureContext({ canvas });
 
@@ -176,17 +175,12 @@ const jumpFlood = root.createGuardedComputePipeline((x, y) => {
   std.textureStore(pingPongLayout.$.writeView, d.vec2i(x, y), 1, d.vec4f(bestSample.coord, 0, 0));
 });
 
-const voronoiFrag = tgpu.fragmentFn({
-  in: { uv: d.vec2f },
-  out: d.vec4f,
-})(({ uv }) =>
-  std.textureSample(colorSampleLayout.$.floodTexture, colorSampleLayout.$.sampler, uv),
-);
-
 const voronoiPipeline = root.createRenderPipeline({
   vertex: common.fullScreenTriangle,
-  fragment: voronoiFrag,
-  targets: { format: presentationFormat },
+  fragment: ({ uv }) => {
+    'use gpu';
+    return std.textureSample(colorSampleLayout.$.floodTexture, colorSampleLayout.$.sampler, uv);
+  },
 });
 
 const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));

apps/typegpu-docs/src/examples/image-processing/background-segmentation/shaders.ts

@@ -1,4 +1,4 @@
-import tgpu, { d, std } from 'typegpu';
+import tgpu, { d, std, type TgpuFragmentFn } from 'typegpu';
 import { MODEL_HEIGHT, MODEL_WIDTH } from './model.ts';
 import {
   blockDim,
@@ -94,9 +94,8 @@ export const computeFn = tgpu.computeFn({
   }
 });
 
-export const fragmentFn = (input: { uv: d.v2f }) => {
+export const fragmentFn = ({ uv }: TgpuFragmentFn.AutoIn<{ uv: d.v2f }>) => {
   'use gpu';
-  const uv = input.uv;
   const originalColor = std.textureSampleBaseClampToEdge(
     drawWithMaskLayout.$.inputTexture,
     drawWithMaskLayout.$.sampler,

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

@@ -1,6 +1,6 @@
 import { linearToSrgb, srgbToLinear } from '@typegpu/color';
-import tgpu, { d } from 'typegpu';
-import { add, discard, div, max, min, mul, normalize, pow, sub } from 'typegpu/std';
+import tgpu, { d, type TgpuFragmentFn, type TgpuVertexFn } from 'typegpu';
+import { discard, max, min, mul, normalize, pow, sub } from 'typegpu/std';
 import { mat4 } from 'wgpu-matrix';
 import { defineControls } from '../../common/defineControls.ts';
 
@@ -133,41 +133,39 @@ const getBoxIntersection = tgpu
 }`)
   .$uses({ IntersectionStruct });
 
-const Varying = {
-  rayWorldOrigin: d.vec3f,
-};
-
-const mainVertex = tgpu.vertexFn({
-  in: { vertexIndex: d.builtin.vertexIndex },
-  out: { pos: d.builtin.position, ...Varying },
-})((input) => {
+const mainVertex = ({ $vertexIndex: vid }: TgpuVertexFn.AutoIn<{}>) => {
+  'use gpu';
   const pos = [d.vec2f(-1, -1), d.vec2f(3, -1), d.vec2f(-1, 3)];
 
-  const rayWorldOrigin = mul(uniforms.$.invViewMatrix, d.vec4f(0, 0, 0, 1)).xyz;
+  const rayWorldOrigin = (uniforms.$.invViewMatrix * d.vec4f(0, 0, 0, 1)).xyz;
 
-  return { pos: d.vec4f(pos[input.vertexIndex], 0.0, 1.0), rayWorldOrigin };
-});
+  return {
+    $position: d.vec4f(pos[vid], 0, 1),
+    rayWorldOrigin,
+  } satisfies TgpuVertexFn.AutoOut;
+};
 
-const fragmentFunction = tgpu.fragmentFn({
-  in: { position: d.builtin.position, ...Varying },
-  out: d.vec4f,
-})((input) => {
-  const boxSize3 = d.vec3f(d.f32(uniforms.$.boxSize));
-  const halfBoxSize3 = mul(0.5, boxSize3);
-  const halfCanvasDims = mul(0.5, uniforms.$.canvasDims);
+const fragmentFunction = ({
+  $position,
+  rayWorldOrigin,
+}: TgpuFragmentFn.AutoIn<{ rayWorldOrigin: d.v3f }>) => {
+  'use gpu';
+  const boxSize3 = d.vec3f(uniforms.$.boxSize);
+  const halfBoxSize3 = 0.5 * boxSize3;
+  const halfCanvasDims = 0.5 * uniforms.$.canvasDims;
 
   const minDim = min(uniforms.$.canvasDims.x, uniforms.$.canvasDims.y);
-  const viewCoords = div(sub(input.position.xy, halfCanvasDims), minDim);
+  const viewCoords = ($position.xy - halfCanvasDims) / minDim;
 
   const ray = Ray({
-    origin: input.rayWorldOrigin,
-    direction: mul(uniforms.$.invViewMatrix, d.vec4f(normalize(d.vec3f(viewCoords, 1)), 0)).xyz,
+    origin: rayWorldOrigin,
+    direction: (uniforms.$.invViewMatrix * d.vec4f(normalize(d.vec3f(viewCoords, 1)), 0)).xyz,
   });
 
   const bigBoxIntersection = getBoxIntersection(
     AxisAlignedBounds({
-      min: mul(-1, halfBoxSize3),
-      max: add(cubeSize, halfBoxSize3),
+      min: -1 * halfBoxSize3,
+      max: cubeSize + halfBoxSize3,
     }),
     ray,
   );
@@ -188,12 +186,12 @@ const fragmentFunction = tgpu.fragmentFn({
           continue;
         }
 
-        const ijkScaled = d.vec3f(d.f32(i), d.f32(j), d.f32(k));
+        const ijkScaled = d.vec3f(i, j, k);
 
         const intersection = getBoxIntersection(
           AxisAlignedBounds({
-            min: sub(ijkScaled, halfBoxSize3),
-            max: add(ijkScaled, halfBoxSize3),
+            min: ijkScaled - halfBoxSize3,
+            max: ijkScaled + halfBoxSize3,
           }),
           ray,
         );
@@ -202,25 +200,25 @@ const fragmentFunction = tgpu.fragmentFn({
           const boxDensity =
             max(0, intersection.tMax - intersection.tMin) * pow(uniforms.$.materialDensity, 2);
           density += boxDensity;
-          invColor = add(invColor, mul(boxDensity, div(d.vec3f(1), boxMatrix.$[i][j][k].albedo)));
+          invColor += boxDensity * (1 / boxMatrix.$[i][j][k].albedo);
           intersectionFound = true;
         }
       }
     }
   }
 
-  const linear = div(d.vec3f(1), invColor);
+  const linear = 1 / invColor;
   const srgb = linearToSrgb(linear);
   const gamma = 2.2;
-  const corrected = pow(srgb, d.vec3f(1.0 / gamma));
+  const corrected = pow(srgb, d.vec3f(1 / gamma));
 
   if (intersectionFound) {
-    return mul(min(density, 1), d.vec4f(min(corrected, d.vec3f(1)), 1));
+    return min(density, 1) * d.vec4f(min(corrected, d.vec3f(1)), 1);
   }
 
   discard();
   return d.vec4f();
-});
+};
 
 // pipeline
 

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

@@ -1,4 +1,4 @@
-import tgpu, { d, std, type AutoVertexIn } from 'typegpu';
+import tgpu, { d, std, type TgpuVertexFn } from 'typegpu';
 import { defineControls } from '../../common/defineControls.ts';
 
 const root = await tgpu.init({
@@ -256,7 +256,7 @@ export const controls = defineControls({
         console.log(n);
       };
 
-      const vs = ({ $vertexIndex }: AutoVertexIn<{}>) => {
+      const vs = ({ $vertexIndex }: TgpuVertexFn.AutoIn<{}>) => {
         'use gpu';
         const positions = [d.vec2f(0, 0.5), d.vec2f(-0.5, -0.5), d.vec2f(0.5, -0.5)];
         myLog(6);

apps/typegpu-docs/tests/individual-example-tests/box-raytracing.test.ts

@@ -29,19 +29,19 @@ describe('box raytracing example', () => {
 
       @group(0) @binding(0) var<uniform> uniforms: Uniforms;
 
-      struct mainVertex_Output {
-        @builtin(position) pos: vec4f,
+      struct VertexOut {
+        @builtin(position) position: vec4f,
         @location(0) rayWorldOrigin: vec3f,
       }
 
-      @vertex fn mainVertex(@builtin(vertex_index) _arg_vertexIndex: u32) -> mainVertex_Output {
-        var pos = array<vec2f, 3>(vec2f(-1), vec2f(3, -1), vec2f(-1, 3));
-        var rayWorldOrigin = (uniforms.invViewMatrix * vec4f(0, 0, 0, 1)).xyz;
-        return mainVertex_Output(vec4f(pos[_arg_vertexIndex], 0f, 1f), rayWorldOrigin);
+      struct VertexIn {
+        @builtin(vertex_index) vertexIndex: u32,
       }
 
-      struct fragmentFunction_Input {
-        @location(0) rayWorldOrigin: vec3f,
+      @vertex fn mainVertex(_arg_0: VertexIn) -> VertexOut {
+        var pos = array<vec2f, 3>(vec2f(-1), vec2f(3, -1), vec2f(-1, 3));
+        var rayWorldOrigin = (uniforms.invViewMatrix * vec4f(0, 0, 0, 1)).xyz;
+        return VertexOut(vec4f(pos[_arg_0.vertexIndex], 0f, 1f), rayWorldOrigin);
       }
 
       struct Ray {
@@ -130,12 +130,17 @@ describe('box raytracing example', () => {
         return select((12.92f * linear), ((1.055f * pow(linear, vec3f(0.4166666567325592))) - vec3f(0.054999999701976776)), (linear > vec3f(0.0031308000907301903)));
       }
 
-      @fragment fn fragmentFunction(_arg_0: fragmentFunction_Input, @builtin(position) _arg_position: vec4f) -> @location(0) vec4f {
+      struct FragmentIn {
+        @builtin(position) position: vec4f,
+        @location(0) rayWorldOrigin: vec3f,
+      }
+
+      @fragment fn fragmentFunction(_arg_0: FragmentIn) -> @location(0) vec4f {
         var boxSize3 = vec3f(uniforms.boxSize);
         var halfBoxSize3 = (0.5f * boxSize3);
         var halfCanvasDims = (0.5f * uniforms.canvasDims);
         let minDim = min(uniforms.canvasDims.x, uniforms.canvasDims.y);
-        var viewCoords = ((_arg_position.xy - halfCanvasDims) / minDim);
+        var viewCoords = ((_arg_0.position.xy - halfCanvasDims) / minDim);
         var ray = Ray(_arg_0.rayWorldOrigin, (uniforms.invViewMatrix * vec4f(normalize(vec3f(viewCoords, 1f)), 0f)).xyz);
         var bigBoxIntersection = getBoxIntersection(AxisAlignedBounds((-1f * halfBoxSize3), (vec3f(7) + halfBoxSize3)), ray);
         if (!bigBoxIntersection.intersects) {
@@ -156,13 +161,13 @@ describe('box raytracing example', () => {
               if (intersection.intersects) {
                 let boxDensity = (max(0f, (intersection.tMax - intersection.tMin)) * pow(uniforms.materialDensity, 2f));
                 density += boxDensity;
-                invColor = (invColor + (boxDensity * (vec3f(1) / boxMatrix[i][j][k].albedo)));
+                invColor += (boxDensity * (1f / boxMatrix[i][j][k].albedo));
                 intersectionFound = true;
               }
             }
           }
         }
-        var linear = (vec3f(1) / invColor);
+        var linear = (1f / invColor);
         var srgb = linearToSrgb(linear);
         const gamma = 2.2;
         var corrected = pow(srgb, vec3f((1f / gamma)));

apps/typegpu-docs/tests/individual-example-tests/jump-flood-voronoi.test.ts

@@ -86,15 +86,15 @@ describe('jump flood (voronoi) example', () => {
         return fullScreenTriangle_Output(vec4f(pos[vertexIndex], 0, 1), uv[vertexIndex]);
       }
 
-      struct voronoiFrag_Input {
-        @location(0) uv: vec2f,
-      }
-
       @group(0) @binding(0) var floodTexture: texture_2d<f32>;
 
       @group(0) @binding(1) var sampler_1: sampler;
 
-      @fragment fn voronoiFrag(_arg_0: voronoiFrag_Input) -> @location(0) vec4f {
+      struct FragmentIn {
+        @location(0) uv: vec2f,
+      }
+
+      @fragment fn fragment(_arg_0: FragmentIn) -> @location(0) vec4f {
         return textureSample(floodTexture, sampler_1, _arg_0.uv);
       }
 

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

@@ -2,10 +2,13 @@ import { builtin, type OmitBuiltins } from '../../builtin.ts';
 import { AutoStruct } from '../../data/autoStruct.ts';
 import type { ResolvedSnippet } from '../../data/snippet.ts';
 import { vec4f } from '../../data/vector.ts';
+import type { FormatToWGSLType } from '../../data/vertexFormatData.ts';
 import type { BaseData, v4f } from '../../data/wgslTypes.ts';
 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 { Assume } from '../../shared/utilityTypes.ts';
+import type { TgpuVertexAttrib } from '../../shared/vertexFormat.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { shaderStageSlot } from '../slot/internalSlots.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
@@ -20,6 +23,12 @@ const builtinVertexIn = {
 
 export type AutoVertexIn<T extends AnyAutoCustoms> = T & InferRecord<typeof builtinVertexIn>;
 
+export type _AutoVertexIn<T> = AutoVertexIn<{
+  [Key in keyof T]: T[Key] extends TgpuVertexAttrib
+    ? InferGPU<FormatToWGSLType<T[Key]['format']>>
+    : Assume<T[Key], InferGPU<BaseIOData>>;
+}>;
+
 const builtinVertexOut = {
   $clipDistances: builtin.clipDistances,
   $position: builtin.position,

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

@@ -22,6 +22,7 @@ import type { Prettify } from '../../shared/utilityTypes.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { addReturnTypeToExternals } from '../resolve/externals.ts';
 import { shaderStageSlot } from '../slot/internalSlots.ts';
+import type { AnyAutoCustoms, AutoFragmentIn, AutoFragmentOut } from './autoIO.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
 import type { BaseIOData, Implementation, InferIO, IOLayout, IORecord } from './fnTypes.ts';
 import { createIoSchema, type IOLayoutToSchema, separateBuiltins } from './ioSchema.ts';
@@ -118,6 +119,8 @@ export declare namespace TgpuFragmentFn {
   // readable, and refactoring to use a builtin argument is too much hassle.
   type In = Record<string, BaseData>;
   type Out = Record<string, BaseData> | BaseData;
+  type AutoIn<T extends AnyAutoCustoms> = AutoFragmentIn<T>;
+  type AutoOut<T extends AnyAutoCustoms = AnyAutoCustoms> = AutoFragmentOut<T>;
 }
 
 export function fragmentFn<FragmentOut extends FragmentOutConstrained>(options: {

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

@@ -7,6 +7,7 @@ import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import type { Prettify } from '../../shared/utilityTypes.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { shaderStageSlot } from '../slot/internalSlots.ts';
+import type { _AutoVertexIn, AnyAutoCustoms, AutoVertexOut } from './autoIO.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
 import type {
   BaseIOData,
@@ -83,6 +84,8 @@ export interface TgpuVertexFn<
 export declare namespace TgpuVertexFn {
   type In = BaseData | Record<string, BaseData>;
   type Out = Record<string, BaseData>;
+  type AutoIn<T> = _AutoVertexIn<T>;
+  type AutoOut<T extends AnyAutoCustoms = AnyAutoCustoms> = AutoVertexOut<T>;
 }
 
 export function vertexFn<VertexOut extends VertexOutConstrained>(options: {

packages/typegpu/src/indexNamedExports.ts

@@ -110,9 +110,13 @@ export type { TgpuVertexFn, TgpuVertexFnShell } from './core/function/tgpuVertex
 export type { TgpuFragmentFn, TgpuFragmentFnShell } from './core/function/tgpuFragmentFn.ts';
 export type { TgpuComputeFn, TgpuComputeFnShell } from './core/function/tgpuComputeFn.ts';
 export type {
+  /** @deprecated use TgpuFragmentFn.AutoIn */
   AutoFragmentIn,
+  /** @deprecated use TgpuFragmentFn.AutoOut */
   AutoFragmentOut,
-  AutoVertexIn,
+  /** @deprecated use TgpuVertexFn.AutoIn */
+  _AutoVertexIn as AutoVertexIn,
+  /** @deprecated use TgpuVertexFn.AutoOut */
   AutoVertexOut,
 } from './core/function/autoIO.ts';
 export type { TgpuDeclare } from './core/declare/tgpuDeclare.ts';