Better DX when creating pipelines

Awesome Blossom

Works maybe?

A few tweaks to make things work a bit more

Author: iwoplaza

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

@@ -169,7 +169,7 @@ const drawWithMaskPipeline = root['~unstable']
   .with(paramsAccess, paramsUniform)
   .createRenderPipeline({
     vertex: common.fullScreenTriangle,
-    fragment: (({ uv }) => {
+    fragment: ({ uv }) => {
       'use gpu';
       const originalColor = std.textureSampleBaseClampToEdge(
         drawWithMaskLayout.$.inputTexture,
@@ -178,7 +178,7 @@ const drawWithMaskPipeline = root['~unstable']
       );
 
       let blurredColor = d.vec4f();
-      if (paramsAccessor.$.useGaussian === 1) {
+      if (paramsAccess.$.useGaussian === 1) {
         blurredColor = std.textureSampleBaseClampToEdge(
           drawWithMaskLayout.$.inputBlurredTexture,
           drawWithMaskLayout.$.sampler,
@@ -189,11 +189,11 @@ const drawWithMaskPipeline = root['~unstable']
           drawWithMaskLayout.$.inputBlurredTexture,
           drawWithMaskLayout.$.sampler,
           uv,
-          paramsAccessor.$.sampleBias,
+          paramsAccess.$.sampleBias,
         );
       }
 
-      const cropBounds = paramsAccessor.$.cropBounds;
+      const cropBounds = paramsAccess.$.cropBounds;
       const uvMin = cropBounds.xy;
       const uvMax = cropBounds.zw;
       const maskUV = d.vec2f(uv).sub(uvMin).div(uvMax.sub(uvMin));
@@ -211,7 +211,7 @@ const drawWithMaskPipeline = root['~unstable']
       const mask = std.select(0, sampledMask, inCropRegion);
 
       return std.mix(blurredColor, originalColor, mask);
-    }),
+    },
     targets: { format: presentationFormat },
   });
 

apps/typegpu-docs/src/examples/simple/oklab/index.ts

@@ -129,7 +129,7 @@ function setPipeline({
   pipeline = root['~unstable']
     .with(patternSlot, outOfGamutPattern)
     .with(oklabGamutClipSlot, gamutClip)
-    .with(oklabGamutClipAlphaAccess, alphaFromUniforms)
+    .with(oklabGamutClipAlphaAccess, () => uniforms.$.alpha)
     .createRenderPipeline({
       vertex: common.fullScreenTriangle,
       fragment: mainFragment,

packages/typegpu/package.json

@@ -65,7 +65,7 @@
   },
   "homepage": "https://typegpu.com",
   "devDependencies": {
-    "@ark/attest": "^0.53.0",
+    "@ark/attest": "^0.56.0",
     "@typegpu/tgpu-dev-cli": "workspace:*",
     "@webgpu/types": "catalog:types",
     "arktype": "catalog:",

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

@@ -1,9 +1,8 @@
-import { builtin } from '../../builtin.ts';
+import { builtin, type OmitBuiltins } from '../../builtin.ts';
 import { AutoStruct } from '../../data/autoStruct.ts';
-import type { AnyData } from '../../data/dataTypes.ts';
 import type { ResolvedSnippet } from '../../data/snippet.ts';
 import { vec4f } from '../../data/vector.ts';
-import type { v4f } from '../../data/wgslTypes.ts';
+import type { BaseData, v4f } from '../../data/wgslTypes.ts';
 import { getName, setName } from '../../shared/meta.ts';
 import type {
   InferGPU,
@@ -35,7 +34,7 @@ const builtinVertexOut = {
 } as const;
 
 export type AutoVertexOut<T extends AnyAutoCustoms> =
-  & T
+  & OmitBuiltins<T>
   & Partial<InferGPURecord<typeof builtinVertexOut>>;
 
 const builtinFragmentIn = {
@@ -62,7 +61,7 @@ export type AutoFragmentOut<T extends undefined | v4f | AnyAutoCustoms> =
     : T & Partial<InferGPURecord<typeof builtinFragmentOut>>;
 
 type AutoFragmentFnImpl = (
-  input: AutoFragmentIn<AnyAutoCustoms>,
+  input: AutoFragmentIn<Record<string, never>>,
 ) => AutoFragmentOut<undefined | v4f | AnyAutoCustoms>;
 
 /**
@@ -80,7 +79,7 @@ export class AutoFragmentFn implements SelfResolvable {
 
   constructor(
     impl: AutoFragmentFnImpl,
-    varyings: Record<string, AnyData>,
+    varyings: Record<string, BaseData>,
     locations?: Record<string, number> | undefined,
   ) {
     this.impl = impl;
@@ -108,7 +107,7 @@ AutoFragmentFn.prototype[$internal] = true;
 AutoFragmentFn.prototype.resourceType = 'auto-fragment-fn';
 
 type AutoVertexFnImpl = (
-  input: AutoVertexIn<AnyAutoCustoms>,
+  input: AutoVertexIn<Record<string, never>>,
 ) => AutoVertexOut<AnyAutoCustoms>;
 
 /**
@@ -126,7 +125,7 @@ export class AutoVertexFn implements SelfResolvable {
 
   constructor(
     impl: AutoVertexFnImpl,
-    attribs: Record<string, AnyData>,
+    attribs: Record<string, BaseData>,
     locations?: Record<string, number> | undefined,
   ) {
     this.impl = impl;

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

@@ -6,32 +6,31 @@ import {
 } from '../../data/attributes.ts';
 import { isBuiltin } from '../../data/attributes.ts';
 import { getCustomLocation, isData } from '../../data/dataTypes.ts';
-import { struct } from '../../data/struct.ts';
+import { INTERNAL_createStruct } from '../../data/struct.ts';
 import {
   type BaseData,
   isVoid,
   type Location,
   type WgslStruct,
 } from '../../data/wgslTypes.ts';
-import type { IOData, IOLayout, IORecord } from './fnTypes.ts';
 
-export type WithLocations<T extends IORecord> = {
+export type WithLocations<T extends Record<string, BaseData>> = {
   [Key in keyof T]: IsBuiltin<T[Key]> extends true ? T[Key]
     : HasCustomLocation<T[Key]> extends true ? T[Key]
     : Decorate<T[Key], Location>;
 };
 
-export type IOLayoutToSchema<T extends IOLayout> = T extends BaseData
-  ? Decorate<T, Location<0>>
-  : T extends IORecord ? WgslStruct<WithLocations<T>>
+export type IOLayoutToSchema<T> = T extends BaseData
+  ? HasCustomLocation<T> extends true ? T : Decorate<T, Location<0>>
+  : T extends Record<string, BaseData> ? WgslStruct<WithLocations<T>>
   // biome-ignore lint/suspicious/noConfusingVoidType: <it actually is void>
   : T extends { type: 'void' } ? void
   : never;
 
-export function withLocations<T extends IOData>(
-  members: IORecord<T> | undefined,
+export function withLocations<T extends BaseData>(
+  members: Record<string, T> | undefined,
   locations: Record<string, number> = {},
-): WithLocations<IORecord<T>> {
+): Record<string, BaseData> {
   let nextLocation = 0;
   const usedCustomLocations = new Set<number>();
 
@@ -69,9 +68,8 @@ export function withLocations<T extends IOData>(
 }
 
 export function createIoSchema<
-  T extends IOData,
-  Layout extends IORecord<T> | IOLayout<T>,
->(layout: Layout, locations: Record<string, number> = {}) {
+  T extends BaseData | Record<string, BaseData>,
+>(layout: T, locations: Record<string, number> = {}) {
   return (
     isData(layout)
       ? isVoid(layout)
@@ -81,6 +79,9 @@ export function createIoSchema<
         : getCustomLocation(layout) !== undefined
         ? layout
         : location(0, layout)
-      : struct(withLocations(layout, locations) as Record<string, T>)
-  ) as IOLayoutToSchema<Layout>;
+      : INTERNAL_createStruct(
+        withLocations(layout as Record<string, BaseData>, locations),
+        /* isAbstruct */ false,
+      )
+  ) as IOLayoutToSchema<T>;
 }

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

@@ -3,12 +3,13 @@ import type {
   AnyFragmentOutputBuiltin,
   OmitBuiltins,
 } from '../../builtin.ts';
+import type { InstanceToSchema } from '../../data/instanceToSchema.ts';
 import type { ResolvedSnippet } from '../../data/snippet.ts';
 import type {
+  BaseData,
   Decorated,
   Interpolate,
   Location,
-  v4f,
   Vec4f,
   Vec4i,
   Vec4u,
@@ -20,7 +21,6 @@ import {
   setName,
   type TgpuNamable,
 } from '../../shared/meta.ts';
-import type { InferGPU } from '../../shared/repr.ts';
 import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { addReturnTypeToExternals } from '../resolve/externals.ts';
@@ -45,6 +45,10 @@ export type FragmentInConstrained = IORecord<
   | AnyFragmentInputBuiltin
 >;
 
+export type FragmentInFromVertexOut<T> =
+  & OmitBuiltins<{ [K in keyof T]: InstanceToSchema<T[K]> }>
+  & Record<string, AnyFragmentInputBuiltin>;
+
 type FragmentColorValue = Vec4f | Vec4i | Vec4u;
 
 export type FragmentOutConstrained = IOLayout<
@@ -53,17 +57,12 @@ export type FragmentOutConstrained = IOLayout<
   | AnyFragmentOutputBuiltin
 >;
 
-export type FragmentOutInferred =
-  | undefined
-  | v4f
-  | Record<string, v4f | InferGPU<AnyFragmentOutputBuiltin>>;
-
 /**
  * Describes a fragment entry function signature (its arguments, return type and targets)
  */
 type TgpuFragmentFnShellHeader<
-  FragmentIn extends FragmentInConstrained,
-  FragmentOut extends FragmentOutConstrained,
+  FragmentIn extends TgpuFragmentFn.In = TgpuFragmentFn.In,
+  FragmentOut extends TgpuFragmentFn.Out = TgpuFragmentFn.Out,
 > = {
   readonly in: FragmentIn | undefined;
   readonly out: FragmentOut;
@@ -76,35 +75,42 @@ type TgpuFragmentFnShellHeader<
  * Allows creating tgpu fragment functions by calling this shell
  * and passing the implementation (as WGSL string or JS function) as the argument.
  */
-export type TgpuFragmentFnShell<
-  FragmentIn extends FragmentInConstrained,
-  FragmentOut extends FragmentOutConstrained,
-> =
-  & TgpuFragmentFnShellHeader<FragmentIn, FragmentOut> /**
+export interface TgpuFragmentFnShell<
+  // We force the variance to be covariant, since shells are just containers of
+  // schemas that conincidentally can be called to create a fragment function.
+  // @ts-expect-error: We override the variance
+  out TIn extends TgpuFragmentFn.In = TgpuFragmentFn.In,
+  // @ts-expect-error: We override the variance
+  out TOut extends TgpuFragmentFn.Out = TgpuFragmentFn.Out,
+> extends TgpuFragmentFnShellHeader<TIn, TOut> {
+  /**
    * Creates a type-safe implementation of this signature
    */
-  & ((
+  (
     implementation: (
-      input: InferIO<FragmentIn>,
-      out: FragmentOut extends IORecord ? WgslStruct<FragmentOut> : FragmentOut,
-    ) => InferIO<FragmentOut>,
-  ) => TgpuFragmentFn<OmitBuiltins<FragmentIn>, FragmentOut>)
-  & /**
+      input: InferIO<TIn>,
+      out: TOut extends IORecord ? WgslStruct<TOut> : TOut,
+    ) => InferIO<TOut>,
+  ): TgpuFragmentFn<OmitBuiltins<TIn>, TOut>;
+  /**
    * @param implementation
    *   Raw WGSL function implementation with header and body
    *   without `fn` keyword and function name
    *   e.g. `"(x: f32) -> f32 { return x; }"`;
-   */ ((
+   */
+  (
     implementation: string,
-  ) => TgpuFragmentFn<OmitBuiltins<FragmentIn>, FragmentOut>)
-  & ((
+  ): TgpuFragmentFn<OmitBuiltins<TIn>, TOut>;
+  (
     strings: TemplateStringsArray,
     ...values: unknown[]
-  ) => TgpuFragmentFn<OmitBuiltins<FragmentIn>, FragmentOut>);
+  ): TgpuFragmentFn<OmitBuiltins<TIn>, TOut>;
+}
 
 export interface TgpuFragmentFn<
-  Varying extends FragmentInConstrained = FragmentInConstrained,
-  Output extends FragmentOutConstrained = FragmentOutConstrained,
+  // @ts-expect-error: We override the variance
+  in Varying extends TgpuFragmentFn.In = Record<string, never>,
+  out Output extends TgpuFragmentFn.Out = TgpuFragmentFn.Out,
 > extends TgpuNamable {
   readonly [$internal]: true;
   readonly shell: TgpuFragmentFnShellHeader<Varying, Output>;
@@ -113,6 +119,13 @@ export interface TgpuFragmentFn<
   $uses(dependencyMap: Record<string, unknown>): this;
 }
 
+export declare namespace TgpuFragmentFn {
+  // Not allowing single-value input, as using objects here is more
+  // readable, and refactoring to use a builtin argument is too much hassle.
+  type In = Record<string, BaseData>;
+  type Out = Record<string, BaseData> | BaseData;
+}
+
 export function fragmentFn<
   FragmentOut extends FragmentOutConstrained,
 >(options: {
@@ -140,10 +153,8 @@ export function fragmentFn<
  *  A `vec4f`, signaling this function outputs a color for one target, or a record containing colors for multiple targets.
  */
 export function fragmentFn<
-  // Not allowing single-value input, as using objects here is more
-  // readable, and refactoring to use a builtin argument is too much hassle.
-  FragmentIn extends FragmentInConstrained,
-  FragmentOut extends FragmentOutConstrained,
+  FragmentIn extends TgpuFragmentFn.In,
+  FragmentOut extends TgpuFragmentFn.Out,
 >(options: {
   in?: FragmentIn;
   out: FragmentOut;
@@ -160,7 +171,7 @@ export function fragmentFn<
     ...values: unknown[]
   ) => createFragmentFn(shell, stripTemplate(arg, ...values));
 
-  return Object.assign(call, shell) as TgpuFragmentFnShell<
+  return Object.assign(call, shell) as unknown as TgpuFragmentFnShell<
     FragmentIn,
     FragmentOut
   >;
@@ -171,16 +182,16 @@ export function fragmentFn<
 // --------------
 
 function createFragmentFn(
-  shell: TgpuFragmentFnShellHeader<
-    FragmentInConstrained,
-    FragmentOutConstrained
-  >,
+  shell: TgpuFragmentFnShellHeader,
   implementation: Implementation,
 ): TgpuFragmentFn {
-  type This = TgpuFragmentFn & SelfResolvable & {
-    [$internal]: true;
-    [$getNameForward]: FnCore;
-  };
+  type This =
+    & TgpuFragmentFn<TgpuFragmentFn.In, TgpuFragmentFn.Out>
+    & SelfResolvable
+    & {
+      [$internal]: true;
+      [$getNameForward]: FnCore;
+    };
 
   const core = createFnCore(implementation, '@fragment ');
   const outputType = shell.returnType;
@@ -214,10 +225,10 @@ function createFragmentFn(
     [$resolve](ctx: ResolutionCtx): ResolvedSnippet {
       const inputWithLocation = shell.in
         ? createIoSchema(shell.in, ctx.varyingLocations)
-          .$name(`${getName(this) ?? ''}_Input`)
         : undefined;
 
       if (inputWithLocation) {
+        setName(inputWithLocation, `${getName(this) ?? ''}_Input`);
         core.applyExternals({ In: inputWithLocation });
       }
       core.applyExternals({ Out: outputType });