Repository: software-mansion/TypeGPU
Pull request #1306 (2025-05-27 11:23:00)

feat: Allow omitting types in wgsl header

github-actions[bot]: **pkg.pr.new**

*packages*
```
pnpm i https://pkg.pr.new/software-mansion/TypeGPU/typegpu@1306
pnpm i https://pkg.pr.new/software-mansion/TypeGPU/typegpu@68a9d30ef058269509e725aa37cd14fa5147e1fd
```

*benchmark*
[view benchmark](https://docs.swmansion.com/TypeGPU/benchmark?p=npm-latest_pr-1306)

*commit*
[view commit](https://github.com/software-mansion/TypeGPU/commit/68a9d30ef058269509e725aa37cd14fa5147e1fd)

Patch:
diff --git a/packages/typegpu/src/core/function/extractArgs.ts b/packages/typegpu/src/core/function/extractArgs.ts
index 4f3ee486..ea6673ba 100644
--- a/packages/typegpu/src/core/function/extractArgs.ts
+++ b/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
diff --git a/packages/typegpu/src/core/function/fnCore.ts b/packages/typegpu/src/core/function/fnCore.ts
index e8a884d0..a95bc89e 100644
--- a/packages/typegpu/src/core/function/fnCore.ts
+++ b/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;
+}


Test patch:
diff --git a/packages/typegpu/tests/extractArgs.test.ts b/packages/typegpu/tests/extractArgs.test.ts
index e6c0a82d..40ea449e 100644
--- a/packages/typegpu/tests/extractArgs.test.ts
+++ b/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 });
+  });
 });
diff --git a/packages/typegpu/tests/rawFn.test.ts b/packages/typegpu/tests/rawFn.test.ts
index d24dbc02..f93dfd03 100644
--- a/packages/typegpu/tests/rawFn.test.ts
+++ b/packages/typegpu/tests/rawFn.test.ts
@@ -342,6 +342,28 @@ struct fragment_Output {
     );
   });
 
+  // TODO: handle nested structs
+  // it('adds return type nested struct definitions when resolving wgsl-defined functions', () => {
+  //   const Point = d.struct({ a: d.u32 }).$name('P');
+
+  //   const func = tgpu['~unstable']
+  //     .fn([d.arrayOf(Point, 4)])(/* wgsl */ `(a: array<MyPoint, 4>) {
+  //       return;
+  //     }`)
+  //     .$name('f');
+
+  //   expect(parseResolved({ func })).toBe(
+  //     parse(`
+  //   struct P {
+  //     a: u32,
+  //   }
+
+  //   fn f(a: array<P, 4>) {
+  //     return;
+  //   }`),
+  //   );
+  // });
+
   it('resolves object externals and replaces their usages in code', () => {
     const getColor = tgpu['~unstable']
       .fn([], d.vec3f)(`() -> vec3f {
@@ -372,6 +394,228 @@ struct fragment_Output {
     `),
     );
   });
+
+  it('resolves compound types with structs provided in externals', () => {
+    const Point = d.struct({ a: d.u32 }).$name('P');
+
+    const getColor = tgpu['~unstable']
+      .fn([d.arrayOf(Point, 4)], d.u32)(
+        /* wgsl */ `(a: array<MyPoint, 4>) {
+        var b: MyPoint = a[0];
+        return b.a;
+      }`,
+      )
+      .$name('get_color')
+      .$uses({ MyPoint: Point });
+
+    expect(parseResolved({ getColor })).toBe(
+      parse(`
+      struct P { a: u32, }
+      fn get_color(a: array<P, 4>) -> u32 {
+        var b: P = a[0];
+        return b.a;
+      }
+    `),
+    );
+  });
+});
+
+describe('tgpu.fn with raw wgsl and missing types', () => {
+  it('resolves missing base types', () => {
+    const getColor = tgpu['~unstable']
+      .fn([d.vec3f, d.u32, d.mat2x2f, d.bool, d.vec2b], d.vec4u)(
+        /* wgsl */ `(a, b: u32, c, d, e) {
+        return vec4u();
+      }`,
+      )
+      .$name('get_color');
+
+    expect(parseResolved({ getColor })).toBe(
+      parse(`
+      fn get_color(a: vec3f, b: u32, c: mat2x2f, d: bool, e: vec2<bool>) -> vec4u {
+        return vec4u();
+      }
+    `),
+    );
+  });
+
+  it('resolves void functions', () => {
+    const getColor = tgpu['~unstable']
+      .fn([])(
+        `() {
+        return;
+      }`,
+      )
+      .$name('get_color');
+
+    expect(parseResolved({ getColor })).toBe(
+      parse(`
+      fn get_color() {
+        return;
+      }
+    `),
+    );
+  });
+
+  it('resolves compound types', () => {
+    const getColor = tgpu['~unstable']
+      .fn([d.arrayOf(d.u32, 4)], d.u32)(
+        /* wgsl */ `(a) {
+        return a[0];
+      }`,
+      )
+      .$name('get_color');
+
+    expect(parseResolved({ getColor })).toBe(
+      parse(`
+      fn get_color(a: array<u32, 4>) -> u32 {
+        return a[0];
+      }
+    `),
+    );
+  });
+
+  it('resolves compound types with structs provided in externals', () => {
+    const Point = d.struct({ a: d.u32 }).$name('P');
+
+    const getColor = tgpu['~unstable']
+      .fn([d.arrayOf(Point, 4)], d.u32)(
+        /* wgsl */ `(a) {
+        var b: MyPoint = a[0];
+        return b.a;
+      }`,
+      )
+      .$name('get_color')
+      .$uses({ MyPoint: Point });
+
+    expect(parseResolved({ getColor })).toBe(
+      parse(`
+      struct P { a: u32, }
+      fn get_color(a: array<P, 4>) -> u32 {
+        var b: P = a[0];
+        return b.a;
+      }
+    `),
+    );
+  });
+
+  it('replaces references when one struct is named in wgsl', () => {
+    const Point = d
+      .struct({
+        a: d.u32,
+        b: d.u32,
+      })
+      .$name('P');
+
+    const func = tgpu['~unstable']
+      .fn([Point, Point], Point)(/* wgsl */ `(a, b: PointStruct) {
+          return b;
+        }`)
+      .$name('newPointF');
+
+    expect(parseResolved({ func })).toBe(
+      parse(`
+        struct P {
+          a: u32,
+          b: u32,
+        }
+
+        fn newPointF(a: P, b: P) -> P {
+          return b;
+        }`),
+    );
+  });
+
+  it('throws when parameter type mismatch', () => {
+    const getColor = tgpu['~unstable']
+      .fn([d.vec3f])(
+        /* wgsl */ `(a: vec4f) {
+        return;
+      }`,
+      )
+      .$name('get_color');
+
+    expect(() => parseResolved({ getColor }))
+      .toThrowErrorMatchingInlineSnapshot(`
+        [Error: Resolution of the following tree failed: 
+        - <root>
+        - fn:get_color: Type mismatch between TGPU shell and WGSL code string: parameter a, JS type "vec3f", WGSL type "vec4f".]
+      `);
+  });
+
+  it('throws when compound parameter type mismatch', () => {
+    const Point = d.struct({ a: d.u32 }).$name('P');
+
+    const getColor = tgpu['~unstable']
+      .fn([d.arrayOf(Point, 4)])(
+        /* wgsl */ `(a: arrayOf<MyPoint, 3>) {
+        return;
+      }`,
+      )
+      .$name('get_color')
+      .$uses({ MyPoint: Point });
+
+    expect(() => parseResolved({ getColor }))
+      .toThrowErrorMatchingInlineSnapshot(`
+        [Error: Resolution of the following tree failed: 
+        - <root>
+        - fn:get_color: Type mismatch between TGPU shell and WGSL code string: parameter a, JS type "array<P,4>", WGSL type "arrayOf<P,3>".]
+      `);
+  });
+
+  it('throws when return type mismatch', () => {
+    const getColor = tgpu['~unstable']
+      .fn([], d.vec4f)(
+        /* wgsl */ `() -> vec2f {
+        return;
+      }`,
+      )
+      .$name('get_color');
+
+    expect(() => parseResolved({ getColor }))
+      .toThrowErrorMatchingInlineSnapshot(`
+        [Error: Resolution of the following tree failed: 
+        - <root>
+        - fn:get_color: Type mismatch between TGPU shell and WGSL code string: return type, JS type "vec4f", WGSL type "vec2f".]
+      `);
+  });
+
+  it('throws when wrong argument count', () => {
+    const getColor = tgpu['~unstable']
+      .fn([d.vec3f, d.vec4f])(
+        /* wgsl */ `(a, b, c) {
+        return;
+      }`,
+      )
+      .$name('get_color');
+
+    expect(() => parseResolved({ getColor }))
+      .toThrowErrorMatchingInlineSnapshot(`
+        [Error: Resolution of the following tree failed: 
+        - <root>
+        - fn:get_color: WGSL implementation has 3 arguments, while the shell has 2 arguments.]
+      `);
+  });
+
+  it('resolves implicitly typed struct without externals', () => {
+    const Point = d.struct({ a: d.i32 }).$name('myStruct');
+    const getColor = tgpu['~unstable']
+      .fn([Point])(
+        /* wgsl */ `(a) {
+        return;
+      }`,
+      )
+      .$name('get_color');
+
+    console.log(parseResolved({ getColor }));
+
+    expect(parseResolved({ getColor })).toBe(parse(`
+      struct myStruct { a: i32, }
+      fn get_color(a: myStruct) {
+        return;
+      }
+      `));
+  });
 });
 
 describe('tgpu.computeFn with raw string WGSL implementation', () => {