feat: More convenient TgpuVertexFn.AutoIn and TgpuVertexFn.AutoOut types

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,5 +1,5 @@
 import { linearToSrgb, srgbToLinear } from '@typegpu/color';
-import tgpu, { d } from 'typegpu';
+import tgpu, { d, type TgpuFragmentFn, type TgpuVertexFn } from 'typegpu';
 import { add, discard, div, 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
 

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';