feat: Allow omitting types in wgsl header (#1306)

Author: aleksanderkatan

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

@@ -145,11 +145,11 @@ function strip(
       };
     }
 
-    if (code.isAt('(') && !argsStart) {
+    if (code.isAt('(') && argsStart === undefined) {
       argsStart = code.pos;
     }
 
-    if (argsStart) {
+    if (argsStart !== undefined) {
       strippedCode += code.str[code.pos];
     }
     code.advanceBy(1); // parsed character

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

@@ -18,6 +18,7 @@ import {
   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> {
@@ -81,31 +82,70 @@ export function createFnCore(
       const id = ctx.names.makeUnique(getName(this));
 
       if (typeof implementation === 'string') {
+        const replacedImpl = replaceExternalsInWgsl(
+          ctx,
+          externalMap,
+          implementation,
+        );
+
         let header = '';
+        let body = '';
 
         if (shell.isEntry) {
-          const input = isWgslStruct(shell.argTypes[0]) ? '(in: In)' : '()';
+          const input = isWgslStruct(shell.argTypes[0])
+            ? `(in: ${ctx.resolve(shell.argTypes[0])})`
+            : '()';
 
           const attributes = isWgslData(shell.returnType)
             ? getAttributesString(shell.returnType)
             : '';
           const output = shell.returnType !== Void
             ? isWgslStruct(shell.returnType)
-              ? '-> Out'
+              ? `-> ${ctx.resolve(shell.returnType)}`
               : `-> ${attributes !== '' ? attributes : '@location(0)'} ${
                 ctx.resolve(shell.returnType)
               }`
             : '';
+
           header = `${input} ${output} `;
+          body = replacedImpl;
+        } else {
+          const providedArgs = extractArgs(replacedImpl);
+
+          if (providedArgs.args.length !== shell.argTypes.length) {
+            throw new Error(
+              `WGSL implementation has ${providedArgs.args.length} arguments, while the shell has ${shell.argTypes.length} arguments.`,
+            );
+          }
+
+          const input = providedArgs.args.map((argInfo, i) =>
+            `${argInfo.identifier}: ${
+              checkAndReturnType(
+                ctx,
+                `parameter ${argInfo.identifier}`,
+                argInfo.type,
+                shell.argTypes[i],
+              )
+            }`
+          ).join(', ');
+
+          const output = shell.returnType === Void
+            ? ''
+            : `-> ${
+              checkAndReturnType(
+                ctx,
+                'return type',
+                providedArgs.ret?.type,
+                shell.returnType,
+              )
+            }`;
+
+          header = `(${input}) ${output}`;
+
+          body = replacedImpl.slice(providedArgs.range.end);
         }
 
-        const replacedImpl = replaceExternalsInWgsl(
-          ctx,
-          externalMap,
-          `${header}${implementation.trim()}`,
-        );
-
-        ctx.addDeclaration(`${fnAttribute}fn ${id}${replacedImpl}`);
+        ctx.addDeclaration(`${fnAttribute}fn ${id}${header}${body}`);
       } else {
         // get data generated by the plugin
         const pluginData = getMetaData(implementation);
@@ -176,3 +216,26 @@ export function createFnCore(
 
   return core;
 }
+
+function checkAndReturnType(
+  ctx: ResolutionCtx,
+  name: string,
+  wgslType: string | undefined,
+  jsType: unknown,
+) {
+  const resolvedJsType = ctx.resolve(jsType).replace(/\s/g, '');
+
+  if (!wgslType) {
+    return resolvedJsType;
+  }
+
+  const resolvedWgslType = wgslType.replace(/\s/g, '');
+
+  if (resolvedWgslType !== resolvedJsType) {
+    throw new Error(
+      `Type mismatch between TGPU shell and WGSL code string: ${name}, JS type "${resolvedJsType}", WGSL type "${resolvedWgslType}".`,
+    );
+  }
+
+  return wgslType;
+}

packages/typegpu/tests/extractArgs.test.ts

@@ -331,4 +331,14 @@ describe('extract args', () => {
     expect(ret).toBeUndefined();
     expect(range).toStrictEqual({ begin: 21, end: 57 });
   });
+
+  it('extracts when no arguments, no name, no return type', () => {
+    const wgslFn = /* wgsl */ '() { return 42; }';
+
+    const { args, ret, range } = extractArgs(wgslFn);
+
+    expect(args).toStrictEqual([]);
+    expect(ret).toStrictEqual(undefined);
+    expect(range).toStrictEqual({ begin: 0, end: 3 });
+  });
 });