Repository: software-mansion/TypeGPU
Pull request #1386 (2025-06-20 08:43:51)

feat: Better TgpuFn type (0.6.0)
Changes:
- Argument names are maintained by the type system (if implemented using TGSL)
- `TgpuFn` without generic instantiation is wide-enough that it represents any function created using `tgpu.fn`

https://github.com/user-attachments/assets/3f79bf10-9d03-4564-8d05-35dd6f5e7248



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

*packages*
```
pnpm i https://pkg.pr.new/software-mansion/TypeGPU/typegpu@1386
pnpm i https://pkg.pr.new/software-mansion/TypeGPU/typegpu@ff61c6af1f76720f06b1ed68885fdc4aeb687f5b
```

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

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

Patch:
diff --git a/apps/typegpu-docs/src/content/examples/rendering/perlin-noise/index.ts b/apps/typegpu-docs/src/content/examples/rendering/perlin-noise/index.ts
index 75e6d1d5..da445350 100644
--- a/apps/typegpu-docs/src/content/examples/rendering/perlin-noise/index.ts
+++ b/apps/typegpu-docs/src/content/examples/rendering/perlin-noise/index.ts
@@ -30,7 +30,9 @@ const tanhSharpen = tgpu['~unstable'].fn([d.f32, d.f32], d.f32)(
   (n, sharpness) => tanh(n * (1 + sharpness * 10)),
 );
 
-const sharpenFnSlot = tgpu['~unstable'].slot<TgpuFn<[d.F32, d.F32], d.F32>>(
+const sharpenFnSlot = tgpu['~unstable'].slot<
+  TgpuFn<(n: d.F32, sharpness: d.F32) => d.F32>
+>(
   exponentialSharpen,
 );
 
diff --git a/apps/typegpu-docs/src/content/examples/simulation/stable-fluid/index.ts b/apps/typegpu-docs/src/content/examples/simulation/stable-fluid/index.ts
index e7299fab..fec98f7e 100644
--- a/apps/typegpu-docs/src/content/examples/simulation/stable-fluid/index.ts
+++ b/apps/typegpu-docs/src/content/examples/simulation/stable-fluid/index.ts
@@ -147,7 +147,7 @@ function createRenderPipeline(
   fragmentFn: TgpuFragmentFn<{ uv: d.Vec2f }, d.Vec4f>,
 ) {
   return root['~unstable']
-    .withVertex(renderFn, renderFn.shell.attributes)
+    .withVertex(renderFn, {})
     .withFragment(fragmentFn, { format })
     .withPrimitive({
       topology: 'triangle-strip',
diff --git a/packages/typegpu-noise/src/generator.ts b/packages/typegpu-noise/src/generator.ts
index a7feba21..18e3d9f2 100644
--- a/packages/typegpu-noise/src/generator.ts
+++ b/packages/typegpu-noise/src/generator.ts
@@ -3,11 +3,11 @@ import * as d from 'typegpu/data';
 import { add, cos, dot, fract } from 'typegpu/std';
 
 export interface StatefulGenerator {
-  seed: TgpuFn<[d.F32], d.Void>;
-  seed2: TgpuFn<[d.Vec2f], d.Void>;
-  seed3: TgpuFn<[d.Vec3f], d.Void>;
-  seed4: TgpuFn<[d.Vec4f], d.Void>;
-  sample: TgpuFn<[], d.F32>;
+  seed: TgpuFn<(seed: d.F32) => d.Void>;
+  seed2: TgpuFn<(seed: d.Vec2f) => d.Void>;
+  seed3: TgpuFn<(seed: d.Vec3f) => d.Void>;
+  seed4: TgpuFn<(seed: d.Vec4f) => d.Void>;
+  sample: TgpuFn<() => d.F32>;
 }
 
 export const randomGeneratorShell: TgpuFnShell<[], d.F32> = tgpu[
diff --git a/packages/typegpu-noise/src/perlin-3d/dynamic-cache.ts b/packages/typegpu-noise/src/perlin-3d/dynamic-cache.ts
index c5fc4bf2..642302af 100644
--- a/packages/typegpu-noise/src/perlin-3d/dynamic-cache.ts
+++ b/packages/typegpu-noise/src/perlin-3d/dynamic-cache.ts
@@ -40,7 +40,7 @@ type Bindings<Prefix extends string> = Prettify<
 export interface DynamicPerlin3DCacheConfig<Prefix extends string> {
   readonly layout: Layout<Prefix>;
   readonly valuesSlot: TgpuSlot<LayoutValue<Prefix>>;
-  readonly getJunctionGradient: TgpuFn<[pos: d.Vec3i], d.Vec3f>;
+  readonly getJunctionGradient: TgpuFn<(pos: d.Vec3i) => d.Vec3f>;
 
   instance(
     root: TgpuRoot,
diff --git a/packages/typegpu-noise/src/perlin-3d/static-cache.ts b/packages/typegpu-noise/src/perlin-3d/static-cache.ts
index 1fe4d51a..e3cf51cf 100644
--- a/packages/typegpu-noise/src/perlin-3d/static-cache.ts
+++ b/packages/typegpu-noise/src/perlin-3d/static-cache.ts
@@ -5,7 +5,7 @@ import { computeJunctionGradient } from './algorithm.ts';
 const MemorySchema = (n: number) => d.arrayOf(d.vec3f, n);
 
 export interface StaticPerlin3DCache {
-  readonly getJunctionGradient: TgpuFn<[pos: d.Vec3i], d.Vec3f>;
+  readonly getJunctionGradient: TgpuFn<(pos: d.Vec3i) => d.Vec3f>;
   readonly size: d.v3u;
   destroy(): void;
 }
diff --git a/packages/typegpu-noise/src/random.ts b/packages/typegpu-noise/src/random.ts
index 809563d6..1b6b4784 100644
--- a/packages/typegpu-noise/src/random.ts
+++ b/packages/typegpu-noise/src/random.ts
@@ -5,81 +5,69 @@ import { randomGeneratorSlot } from './generator.ts';
 
 const TWO_PI = Math.PI * 2;
 
-export const randSeed: TgpuFn<[d.F32], d.Void> = tgpu['~unstable'].fn([d.f32])(
-  (seed) => {
+export const randSeed: TgpuFn<(seed: d.F32) => d.Void> = tgpu['~unstable']
+  .fn([d.f32])((seed) => {
     randomGeneratorSlot.value.seed(seed);
-  },
-);
+  });
 
-export const randSeed2: TgpuFn<[d.Vec2f], d.Void> = tgpu['~unstable'].fn([
-  d.vec2f,
-])((seed) => {
-  randomGeneratorSlot.value.seed2(seed);
-});
+export const randSeed2: TgpuFn<(seed: d.Vec2f) => d.Void> = tgpu['~unstable']
+  .fn([d.vec2f])((seed) => {
+    randomGeneratorSlot.value.seed2(seed);
+  });
 
-export const randSeed3: TgpuFn<[d.Vec3f], d.Void> = tgpu['~unstable'].fn([
-  d.vec3f,
-])((seed) => {
-  randomGeneratorSlot.value.seed3(seed);
-});
+export const randSeed3: TgpuFn<(seed: d.Vec3f) => d.Void> = tgpu['~unstable']
+  .fn([d.vec3f])((seed) => {
+    randomGeneratorSlot.value.seed3(seed);
+  });
 
-export const randSeed4: TgpuFn<[d.Vec4f], d.Void> = tgpu['~unstable'].fn([
-  d.vec4f,
-])((seed) => {
-  randomGeneratorSlot.value.seed4(seed);
-});
+export const randSeed4: TgpuFn<(seed: d.Vec4f) => d.Void> = tgpu['~unstable']
+  .fn([d.vec4f])((seed) => {
+    randomGeneratorSlot.value.seed4(seed);
+  });
 
-export const randFloat01: TgpuFn<[], d.F32> = tgpu['~unstable'].fn([], d.f32)(
+export const randFloat01: TgpuFn<() => d.F32> = tgpu['~unstable'].fn([], d.f32)(
   () => randomGeneratorSlot.value.sample(),
 );
 
-export const randInUnitCube: TgpuFn<[], d.Vec3f> = tgpu['~unstable'].fn(
-  [],
-  d.vec3f,
-)(() =>
-  d.vec3f(
-    randomGeneratorSlot.value.sample() * 2 - 1,
-    randomGeneratorSlot.value.sample() * 2 - 1,
-    randomGeneratorSlot.value.sample() * 2 - 1,
-  )
-);
+export const randInUnitCube: TgpuFn<() => d.Vec3f> = tgpu['~unstable']
+  .fn([], d.vec3f)(() =>
+    d.vec3f(
+      randomGeneratorSlot.value.sample() * 2 - 1,
+      randomGeneratorSlot.value.sample() * 2 - 1,
+      randomGeneratorSlot.value.sample() * 2 - 1,
+    )
+  );
 
-export const randInUnitCircle: TgpuFn<[], d.Vec2f> = tgpu['~unstable'].fn(
-  [],
-  d.vec2f,
-)(() => {
-  const radius = sqrt(randomGeneratorSlot.value.sample());
-  const angle = randomGeneratorSlot.value.sample() * TWO_PI;
+export const randInUnitCircle: TgpuFn<() => d.Vec2f> = tgpu['~unstable']
+  .fn([], d.vec2f)(() => {
+    const radius = sqrt(randomGeneratorSlot.value.sample());
+    const angle = randomGeneratorSlot.value.sample() * TWO_PI;
 
-  return d.vec2f(cos(angle) * radius, sin(angle) * radius);
-});
+    return d.vec2f(cos(angle) * radius, sin(angle) * radius);
+  });
 
-export const randOnUnitCircle: TgpuFn<[], d.Vec2f> = tgpu['~unstable'].fn(
-  [],
-  d.vec2f,
-)(() => {
-  const angle = randomGeneratorSlot.value.sample() * TWO_PI;
-  return d.vec2f(cos(angle), sin(angle));
-});
+export const randOnUnitCircle: TgpuFn<() => d.Vec2f> = tgpu['~unstable']
+  .fn([], d.vec2f)(() => {
+    const angle = randomGeneratorSlot.value.sample() * TWO_PI;
+    return d.vec2f(cos(angle), sin(angle));
+  });
 
-export const randOnUnitSphere: TgpuFn<[], d.Vec3f> = tgpu['~unstable'].fn(
-  [],
-  d.vec3f,
-)(() => {
-  const z = 2 * randomGeneratorSlot.value.sample() - 1;
-  const oneMinusZSq = sqrt(1 - z * z);
-  // TODO: Work out if the -Math.PI offset is necessary
-  const theta = TWO_PI * randomGeneratorSlot.value.sample() - Math.PI;
-  const x = sin(theta) * oneMinusZSq;
-  const y = cos(theta) * oneMinusZSq;
-  return d.vec3f(x, y, z);
-});
+export const randOnUnitSphere: TgpuFn<() => d.Vec3f> = tgpu['~unstable']
+  .fn([], d.vec3f)(() => {
+    const z = 2 * randomGeneratorSlot.value.sample() - 1;
+    const oneMinusZSq = sqrt(1 - z * z);
+    // TODO: Work out if the -Math.PI offset is necessary
+    const theta = TWO_PI * randomGeneratorSlot.value.sample() - Math.PI;
+    const x = sin(theta) * oneMinusZSq;
+    const y = cos(theta) * oneMinusZSq;
+    return d.vec3f(x, y, z);
+  });
 
-export const randOnUnitHemisphere: TgpuFn<[normal: d.Vec3f], d.Vec3f> = tgpu[
-  '~unstable'
-].fn([d.vec3f], d.vec3f)((normal) => {
-  const value = randOnUnitSphere();
-  const alignment = dot(normal, value);
+export const randOnUnitHemisphere: TgpuFn<(normal: d.Vec3f) => d.Vec3f> =
+  tgpu['~unstable']
+    .fn([d.vec3f], d.vec3f)((normal) => {
+      const value = randOnUnitSphere();
+      const alignment = dot(normal, value);
 
-  return mul(sign(alignment), value);
-});
+      return mul(sign(alignment), value);
+    });
diff --git a/packages/typegpu/src/core/function/fnTypes.ts b/packages/typegpu/src/core/function/fnTypes.ts
index 4db3052b..8f457648 100644
--- a/packages/typegpu/src/core/function/fnTypes.ts
+++ b/packages/typegpu/src/core/function/fnTypes.ts
@@ -23,6 +23,8 @@ import type {
 } from '../../data/wgslTypes.ts';
 import type { Infer } from '../../shared/repr.ts';
 
+export type AnyFn = (...args: never[]) => unknown;
+
 /**
  * Information extracted from transpiling a JS function.
  */
@@ -40,15 +42,35 @@ export type InferArgs<T extends unknown[]> = {
   [Idx in keyof T]: Infer<T[Idx]>;
 };
 
-export type JsImplementation<
-  Args extends unknown[] | [] = unknown[] | [],
-  Return = unknown,
-> = (...args: InferArgs<Args>) => Infer<Return>;
+type InheritTupleValues<T, From> = {
+  [K in keyof T]: K extends keyof From ? From[K] : never;
+};
+
+/**
+ * Returns a type that has arg and return types of `T`, but argument
+ * names of `From`
+ *
+ * Wrapped in an object type with `result` prop just so that it's easier
+ * to remove InheritArgNames<...> from Intellisense with Prettify<T>['result']
+ */
+export type InheritArgNames<T extends AnyFn, From extends AnyFn> = {
+  result: (
+    ...args: Parameters<
+      & ((
+        ...args: InheritTupleValues<Parameters<From>, Parameters<T>>
+      ) => ReturnType<T>)
+      & T
+    >
+  ) => ReturnType<T>;
+};
+
+export type InferImplSchema<ImplSchema extends AnyFn> = (
+  ...args: InferArgs<Parameters<ImplSchema>>
+) => Infer<ReturnType<ImplSchema>>;
 
-export type Implementation<
-  Args extends unknown[] | [] = unknown[] | [],
-  Return = unknown,
-> = string | JsImplementation<Args, Return>;
+export type Implementation<ImplSchema extends AnyFn = AnyFn> =
+  | string
+  | InferImplSchema<ImplSchema>;
 
 export type BaseIOData =
   | F32
diff --git a/packages/typegpu/src/core/function/tgpuFn.ts b/packages/typegpu/src/core/function/tgpuFn.ts
index f6b41761..9b787cc4 100644
--- a/packages/typegpu/src/core/function/tgpuFn.ts
+++ b/packages/typegpu/src/core/function/tgpuFn.ts
@@ -26,8 +26,15 @@ import {
   type TgpuSlot,
 } from '../slot/slotTypes.ts';
 import { createFnCore, type FnCore } from './fnCore.ts';
-import type { Implementation, InferArgs, JsImplementation } from './fnTypes.ts';
+import type {
+  AnyFn,
+  Implementation,
+  InferArgs,
+  InferImplSchema,
+  InheritArgNames,
+} from './fnTypes.ts';
 import { stripTemplate } from './templateUtils.ts';
+import type { Prettify } from '../../shared/utilityTypes.ts';
 
 // ----------
 // Public API
@@ -42,7 +49,7 @@ type TgpuFnShellHeader<
 > = {
   readonly [$internal]: true;
   readonly argTypes: Args;
-  readonly returnType: Return | undefined;
+  readonly returnType: Return;
   readonly isEntry: false;
 };
 
@@ -56,14 +63,16 @@ export type TgpuFnShell<
   Return extends AnyData,
 > =
   & TgpuFnShellHeader<Args, Return>
-  & ((
-    implementation: (...args: InferArgs<Args>) => Infer<Return>,
-  ) => TgpuFn<Args, Return>)
-  & ((implementation: string) => TgpuFn<Args, Return>)
+  & (<T extends (...args: InferArgs<Args>) => Infer<Return>>(
+    implementation: T,
+  ) => TgpuFn<
+    Prettify<InheritArgNames<(...args: Args) => Return, T>>['result']
+  >)
+  & ((implementation: string) => TgpuFn<(...args: Args) => Return>)
   & ((
     strings: TemplateStringsArray,
     ...values: unknown[]
-  ) => TgpuFn<Args, Return>)
+  ) => TgpuFn<(...args: Args) => Return>)
   & {
     /**
      * @deprecated Invoke the shell as a function instead.
@@ -73,36 +82,34 @@ export type TgpuFnShell<
         implementation: (
           ...args: InferArgs<Args>
         ) => Infer<Return>,
-      ) => TgpuFn<Args, Return>)
-      & ((implementation: string) => TgpuFn<Args, Return>);
+      ) => TgpuFn<(...args: Args) => Return>)
+      & ((implementation: string) => TgpuFn<(...args: Args) => Return>);
   };
 
-interface TgpuFnBase<
-  Args extends AnyData[],
-  Return extends AnyData,
-> extends TgpuNamable {
+interface TgpuFnBase<ImplSchema extends AnyFn> extends TgpuNamable {
   readonly [$internal]: {
-    implementation: Implementation<Args, Return>;
+    implementation: Implementation<ImplSchema>;
     argTypes: FnArgsConversionHint;
   };
   readonly resourceType: 'function';
-  readonly shell: TgpuFnShellHeader<Args, Return>;
+  readonly shell: TgpuFnShellHeader<
+    Parameters<ImplSchema>,
+    Extract<ReturnType<ImplSchema>, AnyData>
+  >;
   readonly [$providing]?: Providing | undefined;
 
   $uses(dependencyMap: Record<string, unknown>): this;
-  with<T>(slot: TgpuSlot<T>, value: Eventual<T>): TgpuFn<Args, Return>;
+  with<T>(slot: TgpuSlot<T>, value: Eventual<T>): TgpuFn<ImplSchema>;
   with<T extends AnyData>(
     accessor: TgpuAccessor<T>,
-    value: TgpuFn<[], T> | TgpuBufferUsage<T> | Infer<T>,
-  ): TgpuFn<Args, Return>;
+    value: TgpuFn<() => T> | TgpuBufferUsage<T> | Infer<T>,
+  ): TgpuFn<ImplSchema>;
 }
 
-export type TgpuFn<
-  Args extends AnyData[] = AnyData[],
-  Return extends AnyData = AnyData,
-> =
-  & TgpuFnBase<Args, Return>
-  & ((...args: InferArgs<Args>) => Infer<Return>);
+// biome-ignore lint/suspicious/noExplicitAny: the widest type requires `any`
+export type TgpuFn<ImplSchema extends AnyFn = (...args: any[]) => any> =
+  & TgpuFnBase<ImplSchema>
+  & InferImplSchema<ImplSchema>;
 
 export function fn<
   Args extends AnyData[] | [],
@@ -127,18 +134,22 @@ export function fn<
   const call = (
     arg: Implementation | TemplateStringsArray,
     ...values: unknown[]
-  ) => createFn(shell, stripTemplate(arg, ...values));
+  ) =>
+    createFn(
+      shell as unknown as TgpuFnShellHeader<never[], never>,
+      stripTemplate(arg, ...values),
+    );
 
   return Object.assign(Object.assign(call, shell), {
     does: call,
-  }) as TgpuFnShell<Args, Return>;
+  }) as unknown as TgpuFnShell<Args, Return>;
 }
 
 export function isTgpuFn<Args extends AnyData[] | [], Return extends AnyData>(
-  value: unknown | TgpuFn<Args, Return>,
-): value is TgpuFn<Args, Return> {
-  return !!(value as TgpuFn<Args, Return>)?.[$internal] &&
-    (value as TgpuFn<Args, Return>)?.resourceType === 'function';
+  value: unknown | TgpuFn<(...args: Args) => Return>,
+): value is TgpuFn<(...args: Args) => Return> {
+  return !!(value as TgpuFn<(...args: Args) => Return>)?.[$internal] &&
+    (value as TgpuFn<(...args: Args) => Return>)?.resourceType === 'function';
 }
 
 // --------------
@@ -149,11 +160,14 @@ function stringifyPair([slot, value]: SlotValuePair): string {
   return `${getName(slot) ?? '<unnamed>'}=${value}`;
 }
 
-function createFn<Args extends AnyData[], Return extends AnyData>(
-  shell: TgpuFnShellHeader<Args, Return>,
-  implementation: Implementation<Args, Return>,
-): TgpuFn<Args, Return> {
-  type This = TgpuFnBase<Args, Return> & SelfResolvable & {
+function createFn<ImplSchema extends AnyFn>(
+  shell: TgpuFnShellHeader<
+    Parameters<ImplSchema>,
+    Extract<ReturnType<ImplSchema>, AnyData>
+  >,
+  implementation: Implementation<ImplSchema>,
+): TgpuFn<ImplSchema> {
+  type This = TgpuFnBase<ImplSchema> & SelfResolvable & {
     [$getNameForward]: FnCore;
   };
 
@@ -181,7 +195,7 @@ function createFn<Args extends AnyData[], Return extends AnyData>(
     with(
       slot: TgpuSlot<unknown> | TgpuAccessor,
       value: unknown,
-    ): TgpuFn<Args, Return> {
+    ): TgpuFn<ImplSchema> {
       return createBoundFunction(fn, [
         [isAccessor(slot) ? slot.slot : slot, value],
       ]);
@@ -208,7 +222,7 @@ function createFn<Args extends AnyData[], Return extends AnyData>(
     },
   };
 
-  const call = createDualImpl<JsImplementation<Args, Return>>(
+  const call = createDualImpl<InferImplSchema<ImplSchema>>(
     (...args) => {
       if (typeof implementation === 'string') {
         throw new Error(
@@ -230,8 +244,7 @@ function createFn<Args extends AnyData[], Return extends AnyData>(
   call[$internal].implementation = implementation;
 
   const fn = Object.assign(call, fnBase as This) as unknown as TgpuFn<
-    Args,
-    Return
+    ImplSchema
   >;
 
   Object.defineProperty(fn, 'toString', {
@@ -243,12 +256,12 @@ function createFn<Args extends AnyData[], Return extends AnyData>(
   return fn;
 }
 
-function createBoundFunction<Args extends AnyData[], Return extends AnyData>(
-  innerFn: TgpuFn<Args, Return>,
+function createBoundFunction<ImplSchema extends AnyFn>(
+  innerFn: TgpuFn<ImplSchema>,
   pairs: SlotValuePair[],
-): TgpuFn<Args, Return> {
-  type This = TgpuFnBase<Args, Return> & {
-    [$getNameForward]: TgpuFn<Args, Return>;
+): TgpuFn<ImplSchema> {
+  type This = TgpuFnBase<ImplSchema> & {
+    [$getNameForward]: TgpuFn<ImplSchema>;
   };
 
   const fnBase: This = {
@@ -277,7 +290,7 @@ function createBoundFunction<Args extends AnyData[], Return extends AnyData>(
     with(
       slot: TgpuSlot<unknown> | TgpuAccessor,
       value: unknown,
-    ): TgpuFn<Args, Return> {
+    ): TgpuFn<ImplSchema> {
       return createBoundFunction(fn, [
         ...pairs,
         [isAccessor(slot) ? slot.slot : slot, value],
@@ -285,18 +298,18 @@ function createBoundFunction<Args extends AnyData[], Return extends AnyData>(
     },
   };
 
-  const call = createDualImpl(
-    (...args: InferArgs<Args>): unknown => innerFn(...args),
+  const call = createDualImpl<InferImplSchema<ImplSchema>>(
+    (...args) => innerFn(...args),
     (...args) =>
       snip(
         new FnCall(fn, args.map((arg) => arg.value) as Wgsl[]),
         innerFn.shell.returnType ?? UnknownData,
       ),
     'tgpuFnCall',
-    innerFn.shell.argTypes,
+    innerFn.shell.argTypes as AnyData[],
   );
 
-  const fn = Object.assign(call, fnBase) as TgpuFn<Args, Return>;
+  const fn = Object.assign(call, fnBase) as TgpuFn<ImplSchema>;
 
   Object.defineProperty(fn, 'toString', {
     value() {
@@ -311,12 +324,11 @@ function createBoundFunction<Args extends AnyData[], Return extends AnyData>(
   return fn;
 }
 
-class FnCall<Args extends AnyData[], Return extends AnyData>
-  implements SelfResolvable {
-  readonly [$getNameForward]: TgpuFnBase<Args, Return>;
+class FnCall<ImplSchema extends AnyFn> implements SelfResolvable {
+  readonly [$getNameForward]: TgpuFnBase<ImplSchema>;
 
   constructor(
-    private readonly _fn: TgpuFnBase<Args, Return>,
+    private readonly _fn: TgpuFnBase<ImplSchema>,
     private readonly _params: Wgsl[],
   ) {
     this[$getNameForward] = _fn;
diff --git a/packages/typegpu/src/core/root/init.ts b/packages/typegpu/src/core/root/init.ts
index ba47eabe..71a0fcf6 100644
--- a/packages/typegpu/src/core/root/init.ts
+++ b/packages/typegpu/src/core/root/init.ts
@@ -112,7 +112,7 @@ class WithBindingImpl implements WithBinding {
 
   with<T extends AnyWgslData>(
     slot: TgpuSlot<T> | TgpuAccessor<T>,
-    value: T | TgpuFn<[], T> | TgpuBufferUsage<T> | Infer<T>,
+    value: T | TgpuFn<() => T> | TgpuBufferUsage<T> | Infer<T>,
   ): WithBinding {
     return new WithBindingImpl(this._getRoot, [
       ...this._slotBindings,
diff --git a/packages/typegpu/src/core/root/rootTypes.ts b/packages/typegpu/src/core/root/rootTypes.ts
index cdf815e9..b70c2fa1 100644
--- a/packages/typegpu/src/core/root/rootTypes.ts
+++ b/packages/typegpu/src/core/root/rootTypes.ts
@@ -115,7 +115,7 @@ export interface WithBinding {
   with<T>(slot: TgpuSlot<T>, value: Eventual<T>): WithBinding;
   with<T extends AnyWgslData>(
     accessor: TgpuAccessor<T>,
-    value: TgpuFn<[], T> | TgpuBufferUsage<T> | Infer<T>,
+    value: TgpuFn<() => T> | TgpuBufferUsage<T> | Infer<T>,
   ): WithBinding;
 
   withCompute<ComputeIn extends IORecord<AnyComputeBuiltin>>(
diff --git a/packages/typegpu/src/core/slot/accessor.ts b/packages/typegpu/src/core/slot/accessor.ts
index 77f13f37..29224612 100644
--- a/packages/typegpu/src/core/slot/accessor.ts
+++ b/packages/typegpu/src/core/slot/accessor.ts
@@ -26,7 +26,7 @@ import type { TgpuAccessor, TgpuSlot } from './slotTypes.ts';
 
 export function accessor<T extends AnyWgslData>(
   schema: T,
-  defaultValue?: TgpuFn<[], T> | TgpuBufferUsage<T> | Infer<T>,
+  defaultValue?: TgpuFn<() => T> | TgpuBufferUsage<T> | Infer<T>,
 ): TgpuAccessor<T> {
   return new TgpuAccessorImpl(schema, defaultValue);
 }
@@ -38,7 +38,9 @@ export function accessor<T extends AnyWgslData>(
 export class TgpuAccessorImpl<T extends AnyWgslData>
   implements TgpuAccessor<T>, SelfResolvable {
   public readonly resourceType = 'accessor';
-  public readonly slot: TgpuSlot<TgpuFn<[], T> | TgpuBufferUsage<T> | Infer<T>>;
+  public readonly slot: TgpuSlot<
+    TgpuFn<() => T> | TgpuBufferUsage<T> | Infer<T>
+  >;
 
   // Type-tokens, not available at runtime
   declare readonly [$repr]: Infer<T>;
@@ -49,7 +51,7 @@ export class TgpuAccessorImpl<T extends AnyWgslData>
   constructor(
     public readonly schema: T,
     public readonly defaultValue:
-      | TgpuFn<[], T>
+      | TgpuFn<() => T>
       | TgpuBufferUsage<T>
       | Infer<T>
       | undefined = undefined,
diff --git a/packages/typegpu/src/core/slot/slotTypes.ts b/packages/typegpu/src/core/slot/slotTypes.ts
index e86f1c81..0e51dd1f 100644
--- a/packages/typegpu/src/core/slot/slotTypes.ts
+++ b/packages/typegpu/src/core/slot/slotTypes.ts
@@ -52,11 +52,11 @@ export interface TgpuAccessor<T extends AnyData = AnyData> extends TgpuNamable {
 
   readonly schema: T;
   readonly defaultValue:
-    | TgpuFn<[], T>
+    | TgpuFn<() => T>
     | TgpuBufferUsage<T>
     | Infer<T>
     | undefined;
-  readonly slot: TgpuSlot<TgpuFn<[], T> | TgpuBufferUsage<T> | Infer<T>>;
+  readonly slot: TgpuSlot<TgpuFn<() => T> | TgpuBufferUsage<T> | Infer<T>>;
 
   readonly value: InferGPU<T>;
 }
diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts
index 7645b61d..b3cb450f 100644
--- a/packages/typegpu/src/types.ts
+++ b/packages/typegpu/src/types.ts
@@ -65,8 +65,7 @@ export type ResolvableObject =
   | AnyVecInstance
   | AnyMatInstance
   | AnyData
-  // biome-ignore lint/suspicious/noExplicitAny: <has to be more permissive than unknown>
-  | TgpuFn<any, any>;
+  | TgpuFn;
 
 export type Wgsl = Eventual<string | number | boolean | ResolvableObject>;
 


Test patch:
diff --git a/packages/typegpu/setupVitest.ts b/packages/typegpu/setupVitest.ts
index 497f29d5..f180b61a 100644
--- a/packages/typegpu/setupVitest.ts
+++ b/packages/typegpu/setupVitest.ts
@@ -2,5 +2,5 @@ import { setup } from '@ark/attest';
 
 export default () =>
   setup({
-    formatCmd: 'pnpm check',
+    formatCmd: 'pnpm fix',
   });
diff --git a/packages/typegpu/tests/function.test.ts b/packages/typegpu/tests/function.test.ts
index 8195f9e0..a221dee7 100644
--- a/packages/typegpu/tests/function.test.ts
+++ b/packages/typegpu/tests/function.test.ts
@@ -1,9 +1,15 @@
+import { attest } from '@ark/attest';
 import { describe, expect, expectTypeOf, it } from 'vitest';
-import type { InferIO, IOLayout } from '../src/core/function/fnTypes.ts';
+import type {
+  InferIO,
+  InheritArgNames,
+  IOLayout,
+} from '../src/core/function/fnTypes.ts';
 import * as d from '../src/data/index.ts';
 import tgpu, { type TgpuFn, type TgpuFnShell } from '../src/index.ts';
 import { parse, parseResolved } from './utils/parseResolved.ts';
 import { Void } from '../src/data/wgslTypes.ts';
+import type { Prettify } from '../src/shared/utilityTypes.ts';
 
 describe('tgpu.fn', () => {
   it('should inject function declaration of called function', () => {
@@ -107,16 +113,18 @@ describe('tgpu.fn', () => {
     const two = tgpu['~unstable'].fn([d.f32, d.u32]);
 
     expectTypeOf(proc).toEqualTypeOf<TgpuFnShell<[], d.Void>>();
-    expectTypeOf<ReturnType<typeof proc>>().toEqualTypeOf<TgpuFn<[], d.Void>>();
+    expectTypeOf<ReturnType<typeof proc>>().toEqualTypeOf<
+      TgpuFn<() => d.Void>
+    >();
 
     expectTypeOf(one).toEqualTypeOf<TgpuFnShell<[d.F32], d.Void>>();
     expectTypeOf<ReturnType<typeof one>>().toEqualTypeOf<
-      TgpuFn<[d.F32], d.Void>
+      TgpuFn<(arg_0: d.F32) => d.Void>
     >();
 
     expectTypeOf(two).toEqualTypeOf<TgpuFnShell<[d.F32, d.U32], d.Void>>();
     expectTypeOf<ReturnType<typeof two>>().toEqualTypeOf<
-      TgpuFn<[d.F32, d.U32], d.Void>
+      TgpuFn<(arg_0: d.F32, arg_1: d.U32) => d.Void>
     >();
   });
 
@@ -126,16 +134,18 @@ describe('tgpu.fn', () => {
     const two = tgpu['~unstable'].fn([d.f32, d.u32], d.bool);
 
     expectTypeOf(proc).toEqualTypeOf<TgpuFnShell<[], d.Bool>>();
-    expectTypeOf<ReturnType<typeof proc>>().toEqualTypeOf<TgpuFn<[], d.Bool>>();
+    expectTypeOf<ReturnType<typeof proc>>().toEqualTypeOf<
+      TgpuFn<() => d.Bool>
+    >();
 
     expectTypeOf(one).toEqualTypeOf<TgpuFnShell<[d.F32], d.Bool>>();
     expectTypeOf<ReturnType<typeof one>>().toEqualTypeOf<
-      TgpuFn<[d.F32], d.Bool>
+      TgpuFn<(arg_0: d.F32) => d.Bool>
     >();
 
     expectTypeOf(two).toEqualTypeOf<TgpuFnShell<[d.F32, d.U32], d.Bool>>();
     expectTypeOf<ReturnType<typeof two>>().toEqualTypeOf<
-      TgpuFn<[d.F32, d.U32], d.Bool>
+      TgpuFn<(arg_0: d.F32, arg_1: d.U32) => d.Bool>
     >();
   });
 });
@@ -187,20 +197,6 @@ describe('tgpu.vertexFn', () => {
     expect(parseResolved({ foo })).toContain(parse('struct foo_Out'));
     expect(foo.shell.argTypes).toStrictEqual([]);
   });
-
-  it('does not create In struct when there is empty object for arguments', () => {
-    const foo = tgpu['~unstable'].vertexFn({
-      in: {},
-      out: { pos: d.builtin.position },
-    })(() => {
-      return {
-        pos: d.vec4f(),
-      };
-    });
-    expect(parseResolved({ foo })).not.toContain(parse('struct foo_In'));
-    expect(parseResolved({ foo })).toContain(parse('struct foo_Out'));
-    expect(foo.shell.argTypes).toStrictEqual([]);
-  });
 });
 
 describe('tgpu.fragmentFn', () => {
@@ -245,3 +241,21 @@ describe('InferIO', () => {
     }>();
   });
 });
+
+describe('InheritArgNames', () => {
+  it('should inherit argument names from one fn to another', () => {
+    const isEven = (x: number) => (x & 1) === 0;
+    const identity = (num: number) => num;
+    // Should have the same argument names as `identity`, but the signature of `isEven`
+    const isEvenWithNames = undefined as unknown as Prettify<
+      InheritArgNames<
+        typeof isEven,
+        typeof identity
+      >
+    >['result'];
+
+    attest(isEven).type.toString.snap('(x: number) => boolean');
+    attest(identity).type.toString.snap('(num: number) => number');
+    attest(isEvenWithNames).type.toString.snap('(num: number) => boolean');
+  });
+});
diff --git a/packages/typegpu/tests/tgslFn.test.ts b/packages/typegpu/tests/tgslFn.test.ts
index 17c3adba..38b056d2 100644
--- a/packages/typegpu/tests/tgslFn.test.ts
+++ b/packages/typegpu/tests/tgslFn.test.ts
@@ -4,6 +4,7 @@ import * as d from '../src/data/index.ts';
 import tgpu from '../src/index.ts';
 import { getName } from '../src/shared/meta.ts';
 import { parse, parseResolved } from './utils/parseResolved.ts';
+import { attest } from '@ark/attest';
 
 describe('TGSL tgpu.fn function', () => {
   it('is namable', () => {
@@ -836,4 +837,20 @@ describe('TGSL tgpu.fn function', () => {
 
     expect(actual).toBe(expected);
   });
+
+  it('maintains argument names in the type', () => {
+    const fun = tgpu['~unstable'].fn([d.f32, d.f32], d.f32)((x, y) => {
+      return x + y;
+    });
+
+    attest(fun).type.toString.snap('TgpuFn<(x: F32, y: F32) => F32>');
+  });
+
+  it('falls back to args_N naming when not every argument is used in the implementation', () => {
+    const fun = tgpu['~unstable'].fn([d.f32, d.f32], d.f32)((x) => {
+      return x * 2;
+    });
+
+    attest(fun).type.toString.snap('TgpuFn<(args_0: F32, args_1: F32) => F32>');
+  });
 });