feat: Allow omitting types in wgsl header

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

@@ -16,6 +16,7 @@ import {
   replaceExternalsInWgsl,
 } from '../resolve/externals.ts';
 import { getPrebuiltAstFor } from './astUtils.ts';
+import { extractArgs } from './extractArgs.ts';
 import type { Implementation } from './fnTypes.ts';
 
 export interface TgpuFnShellBase<Args extends unknown[], Return> {
@@ -80,6 +81,7 @@ export function createFnCore(
 
       if (typeof implementation === 'string') {
         let header = '';
+        let body = '';
 
         if (shell.isEntry) {
           const input = isWgslStruct(shell.argTypes[0]) ? '(in: In)' : '()';
@@ -95,12 +97,35 @@ export function createFnCore(
               }`
             : '';
           header = `${input} ${output} `;
+          body = implementation;
+        } else {
+          const providedArgs = extractArgs(implementation);
+
+          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}: ${
+              argInfo.type ?? ctx.resolve(shell.argTypes[i])
+            }`
+          ).join(', ');
+
+          const output = shell.returnType === Void
+            ? ''
+            : `-> ${providedArgs.ret?.type ?? ctx.resolve(shell.returnType)}`;
+
+          header = `(${input}) ${output}`;
+
+          body = implementation.slice(providedArgs.range.end);
         }
 
         const replacedImpl = replaceExternalsInWgsl(
           ctx,
           externalMap,
-          `${header}${implementation.trim()}`,
+          `${header}${body.trim()}`,
         );
 
         ctx.addDeclaration(`${fnAttribute}fn ${id}${replacedImpl}`);

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 });
+  });
 });

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 {
@@ -374,6 +396,151 @@ struct fragment_Output {
   });
 });
 
+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;
+      }
+    `),
+    );
+  });
+
+  // TODO: handle nested structs, which blocks type checking
+  // it('throws when parameter type mismatch', () => {
+  //   const getColor = tgpu['~unstable']
+  //     .fn([d.vec3f])(
+  //       /* wgsl */ `(a: vec4f) {
+  //       return;
+  //     }`,
+  //     )
+  //     .$name('get_color');
+
+  //   expect(() => parseResolved({ getColor }))
+  //     .toThrowErrorMatchingInlineSnapshot();
+  // });
+
+  // TODO: handle nested structs, which blocks type checking
+  // it('throws when return type mismatch', () => {
+  //   const getColor = tgpu['~unstable']
+  //     .fn([], d.vec4f)(
+  //       /* wgsl */ `() -> vec2f {
+  //       return;
+  //     }`,
+  //     )
+  //     .$name('get_color');
+
+  //   expect(() => parseResolved({ getColor }))
+  //     .toThrowErrorMatchingInlineSnapshot();
+  // });
+
+  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', () => {
   it('does not replace supposed input arg types in code', () => {
     const foo = tgpu['~unstable'].computeFn({