feat: Match up vertex and fragment locations in render pipeline

apps/typegpu-docs/src/content/examples/simulation/fluid-with-atomics/index.ts

@@ -285,7 +285,7 @@ const vertex = tgpu['~unstable'].vertexFn({
 
 const fragment = tgpu['~unstable'].fragmentFn({
   in: { cell: d.f32 },
-  out: d.location(0, d.vec4f),
+  out: d.vec4f,
 })((input) => {
   if (input.cell === -1) {
     return d.vec4f(0.5, 0.5, 0.5, 1);

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

@@ -1,8 +1,7 @@
 import { FuncParameterType } from 'tinyest';
 import { getAttributesString } from '../../data/attributes.ts';
-import { snip } from '../../data/dataTypes.ts';
+import { type AnyData, snip } from '../../data/dataTypes.ts';
 import {
-  type AnyWgslData,
   type AnyWgslStruct,
   isWgslData,
   isWgslStruct,
@@ -12,29 +11,25 @@ import { MissingLinksError } from '../../errors.ts';
 import { getMetaData, getName, setName } from '../../shared/meta.ts';
 import type { ResolutionCtx } from '../../types.ts';
 import {
-  addArgTypesToExternals,
-  addReturnTypeToExternals,
   applyExternals,
   type ExternalMap,
   replaceExternalsInWgsl,
 } from '../resolve/externals.ts';
 import { extractArgs } from './extractArgs.ts';
 import type { Implementation } from './fnTypes.ts';
 
-export interface TgpuFnShellBase<Args extends unknown[], Return> {
-  readonly argTypes: Args;
-  readonly returnType: Return;
-  readonly isEntry: boolean;
-}
-
 export interface FnCore {
   applyExternals(newExternals: ExternalMap): void;
-  resolve(ctx: ResolutionCtx, fnAttribute?: string): string;
+  resolve(
+    ctx: ResolutionCtx,
+    argTypes: AnyData[],
+    returnType: AnyData,
+  ): string;
 }
 
 export function createFnCore(
-  shell: TgpuFnShellBase<unknown[], unknown>,
   implementation: Implementation,
+  fnAttribute = '',
 ): FnCore {
   /**
    * External application has to be deferred until resolution because
@@ -44,35 +39,16 @@ export function createFnCore(
    */
   const externalsToApply: ExternalMap[] = [];
 
-  if (typeof implementation === 'string') {
-    if (!shell.isEntry) {
-      addArgTypesToExternals(
-        implementation,
-        shell.argTypes,
-        (externals) => externalsToApply.push(externals),
-      );
-      addReturnTypeToExternals(
-        implementation,
-        shell.returnType,
-        (externals) => externalsToApply.push(externals),
-      );
-    } else {
-      if (isWgslStruct(shell.argTypes[0])) {
-        externalsToApply.push({ In: shell.argTypes[0] });
-      }
-
-      if (isWgslStruct(shell.returnType)) {
-        externalsToApply.push({ Out: shell.returnType });
-      }
-    }
-  }
-
   const core = {
     applyExternals(newExternals: ExternalMap): void {
       externalsToApply.push(newExternals);
     },
 
-    resolve(ctx: ResolutionCtx, fnAttribute = ''): string {
+    resolve(
+      ctx: ResolutionCtx,
+      argTypes: AnyData[],
+      returnType: AnyData,
+    ): string {
       const externalMap: ExternalMap = {};
 
       for (const externals of externalsToApply) {
@@ -91,19 +67,19 @@ export function createFnCore(
         let header = '';
         let body = '';
 
-        if (shell.isEntry) {
-          const input = isWgslStruct(shell.argTypes[0])
-            ? `(in: ${ctx.resolve(shell.argTypes[0])})`
+        if (fnAttribute !== '') {
+          const input = isWgslStruct(argTypes[0])
+            ? `(in: ${ctx.resolve(argTypes[0])})`
             : '()';
 
-          const attributes = isWgslData(shell.returnType)
-            ? getAttributesString(shell.returnType)
+          const attributes = isWgslData(returnType)
+            ? getAttributesString(returnType)
             : '';
-          const output = shell.returnType !== Void
-            ? isWgslStruct(shell.returnType)
-              ? `-> ${ctx.resolve(shell.returnType)}`
+          const output = returnType !== Void
+            ? isWgslStruct(returnType)
+              ? `-> ${ctx.resolve(returnType)}`
               : `-> ${attributes !== '' ? attributes : '@location(0)'} ${
-                ctx.resolve(shell.returnType)
+                ctx.resolve(returnType)
               }`
             : '';
 
@@ -112,9 +88,9 @@ export function createFnCore(
         } else {
           const providedArgs = extractArgs(replacedImpl);
 
-          if (providedArgs.args.length !== shell.argTypes.length) {
+          if (providedArgs.args.length !== argTypes.length) {
             throw new Error(
-              `WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${shell.argTypes.length} arguments.`,
+              `WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${argTypes.length} arguments.`,
             );
           }
 
@@ -124,21 +100,19 @@ export function createFnCore(
                 ctx,
                 `parameter ${argInfo.identifier}`,
                 argInfo.type,
-                shell.argTypes[i],
+                argTypes[i],
               )
             }`
           ).join(', ');
 
-          const output = shell.returnType === Void
-            ? ''
-            : `-> ${
-              checkAndReturnType(
-                ctx,
-                'return type',
-                providedArgs.ret?.type,
-                shell.returnType,
-              )
-            }`;
+          const output = returnType === Void ? '' : `-> ${
+            checkAndReturnType(
+              ctx,
+              'return type',
+              providedArgs.ret?.type,
+              returnType,
+            )
+          }`;
 
           header = `(${input}) ${output}`;
 
@@ -177,12 +151,12 @@ export function createFnCore(
 
         // generate wgsl string
         const { head, body } = ctx.fnToWgsl({
-          args: shell.argTypes.map((arg, i) =>
+          args: argTypes.map((arg, i) =>
             snip(
               ast.params[i]?.type === FuncParameterType.identifier
                 ? ast.params[i].name
                 : `_arg_${i}`,
-              arg as AnyWgslData,
+              arg,
             )
           ),
           argAliases: Object.fromEntries(
@@ -192,14 +166,14 @@ export function createFnCore(
                   alias,
                   snip(
                     `_arg_${i}.${name}`,
-                    (shell.argTypes[i] as AnyWgslStruct)
-                      .propTypes[name] as AnyWgslData,
+                    (argTypes[i] as AnyWgslStruct)
+                      .propTypes[name],
                   ),
                 ])
                 : []
             ),
           ),
-          returnType: shell.returnType as AnyWgslData,
+          returnType,
           body: ast.body,
           externalMap,
         });

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

@@ -29,24 +29,40 @@ export type IOLayoutToSchema<T extends IOLayout> = T extends BaseData
   : never;
 
 export function withLocations<T extends IOData>(
-  members: IORecord<T>,
+  members: IORecord<T> | undefined,
+  locations: Record<string, number> = {},
 ): WithLocations<IORecord<T>> {
   let nextLocation = 0;
+  const usedCustomLocations = new Set<number>();
 
   return Object.fromEntries(
-    Object.entries(members).map(([key, member]) => {
-      if (isBuiltin(member)) {
-        // Skipping builtins
+    Object.entries(members ?? {}).map(([key, member]) => {
+      const customLocation = getCustomLocation(member);
+
+      if (customLocation !== undefined) {
+        if (usedCustomLocations.has(customLocation)) {
+          throw new Error('Duplicate custom location attributes found');
+        }
+        usedCustomLocations.add(customLocation);
+      }
+
+      return [key, member] as const;
+    }).map(([key, member]) => {
+      if (isBuiltin(member)) { // skipping builtins
         return [key, member];
       }
 
-      const customLocation = getCustomLocation(member);
-      if (customLocation !== undefined) {
-        // This member is already marked, start counting from the next location over.
-        nextLocation = customLocation + 1;
+      if (getCustomLocation(member) !== undefined) { // this member is already marked
         return [key, member];
       }
 
+      if (locations[key]) { // location has been determined by a previous procedure
+        return [key, location(locations[key], member)];
+      }
+
+      while (usedCustomLocations.has(nextLocation)) {
+        nextLocation++;
+      }
       return [key, location(nextLocation++, member)];
     }),
   );
@@ -55,14 +71,14 @@ export function withLocations<T extends IOData>(
 export function createIoSchema<
   T extends IOData,
   Layout extends IORecord<T> | IOLayout<T>,
->(returnType: Layout) {
+>(layout: Layout, locations: Record<string, number> = {}) {
   return (
-    isData(returnType)
-      ? isVoid(returnType)
-        ? returnType
-        : getCustomLocation(returnType) !== undefined
-        ? returnType
-        : location(0, returnType)
-      : struct(withLocations(returnType) as Record<string, T>)
+    isData(layout)
+      ? isVoid(layout)
+        ? layout
+        : getCustomLocation(layout) !== undefined
+        ? layout
+        : location(0, layout)
+      : struct(withLocations(layout, locations) as Record<string, T>)
   ) as IOLayoutToSchema<Layout>;
 }

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

@@ -10,7 +10,7 @@ import { $getNameForward, $internal } from '../../shared/symbols.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
 import type { Implementation, InferIO, IORecord } from './fnTypes.ts';
-import { createIoSchema, type IOLayoutToSchema } from './ioOutputType.ts';
+import { createIoSchema, type IOLayoutToSchema } from './ioSchema.ts';
 import { stripTemplate } from './templateUtils.ts';
 
 // ----------
@@ -153,7 +153,10 @@ function createComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
     [$getNameForward]: FnCore;
   };
 
-  const core = createFnCore(shell, implementation);
+  const core = createFnCore(
+    implementation,
+    `@compute @workgroup_size(${workgroupSize.join(', ')}) `,
+  );
   const inputType = shell.argTypes[0];
 
   const result: This = {
@@ -177,7 +180,8 @@ function createComputeFn<ComputeIn extends IORecord<AnyComputeBuiltin>>(
     '~resolve'(ctx: ResolutionCtx): string {
       return core.resolve(
         ctx,
-        `@compute @workgroup_size(${workgroupSize.join(', ')}) `,
+        shell.argTypes,
+        shell.returnType,
       );
     },
 

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

@@ -1,14 +1,15 @@
 import { type AnyData, snip, UnknownData } from '../../data/dataTypes.ts';
 import { Void } from '../../data/wgslTypes.ts';
+import { createDualImpl } from '../../shared/generators.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
 import { getName, setName } from '../../shared/meta.ts';
-import { createDualImpl } from '../../shared/generators.ts';
 import type { Infer } from '../../shared/repr.ts';
 import {
   $getNameForward,
   $internal,
   $providing,
 } from '../../shared/symbols.ts';
+import type { Prettify } from '../../shared/utilityTypes.ts';
 import type { GenerationCtx } from '../../tgsl/generationHelpers.ts';
 import type {
   FnArgsConversionHint,
@@ -17,6 +18,10 @@ import type {
   Wgsl,
 } from '../../types.ts';
 import type { TgpuBufferUsage } from '../buffer/bufferUsage.ts';
+import {
+  addArgTypesToExternals,
+  addReturnTypeToExternals,
+} from '../resolve/externals.ts';
 import {
   type Eventual,
   isAccessor,
@@ -34,7 +39,6 @@ import type {
   InheritArgNames,
 } from './fnTypes.ts';
 import { stripTemplate } from './templateUtils.ts';
-import type { Prettify } from '../../shared/utilityTypes.ts';
 
 // ----------
 // Public API
@@ -157,7 +161,7 @@ function createFn<ImplSchema extends AnyFn>(
     [$getNameForward]: FnCore;
   };
 
-  const core = createFnCore(shell, implementation as Implementation);
+  const core = createFnCore(implementation as Implementation, '');
 
   const fnBase: This = {
     [$internal]: {
@@ -189,7 +193,18 @@ function createFn<ImplSchema extends AnyFn>(
 
     '~resolve'(ctx: ResolutionCtx): string {
       if (typeof implementation === 'string') {
-        return core.resolve(ctx);
+        addArgTypesToExternals(
+          implementation,
+          shell.argTypes,
+          core.applyExternals,
+        );
+        addReturnTypeToExternals(
+          implementation,
+          shell.returnType,
+          core.applyExternals,
+        );
+
+        return core.resolve(ctx, shell.argTypes, shell.returnType);
       }
 
       const generationCtx = ctx as GenerationCtx;
@@ -201,7 +216,7 @@ function createFn<ImplSchema extends AnyFn>(
 
       try {
         generationCtx.callStack.push(shell.returnType);
-        return core.resolve(ctx);
+        return core.resolve(ctx, shell.argTypes, shell.returnType);
       } finally {
         generationCtx.callStack.pop();
       }