From a0af2518758b1924300ccc379f68dc3131ab54de Mon Sep 17 00:00:00 2001
From: Iwo Plaza <iwoplaza@gmail.com>
Date: Thu, 21 May 2026 16:07:10 +0200
Subject: [PATCH] feat: std.typeOf

---
 .../src/content/docs/apis/utils.mdx           |   3 +
 packages/typegpu-testing-utility/src/index.ts |  32 ++
 packages/typegpu/src/data/index.ts            |   4 +-
 packages/typegpu/src/std/index.ts             |   2 +
 packages/typegpu/src/std/typeOf.ts            |  68 ++++
 packages/typegpu/tests/std/typeOf.test.ts     | 354 ++++++++++++++++++
 6 files changed, 462 insertions(+), 1 deletion(-)
 create mode 100644 packages/typegpu/src/std/typeOf.ts
 create mode 100644 packages/typegpu/tests/std/typeOf.test.ts

diff --git a/apps/typegpu-docs/src/content/docs/apis/utils.mdx b/apps/typegpu-docs/src/content/docs/apis/utils.mdx
index 7af6f5c7a3..f6d70a193b 100644
--- a/apps/typegpu-docs/src/content/docs/apis/utils.mdx
+++ b/apps/typegpu-docs/src/content/docs/apis/utils.mdx
@@ -118,6 +118,9 @@ TypeGPU exposes standard-library probes for code that needs to adapt to its exec
   if it's called during shader generation, and `undefined` otherwise.
 - `std.getShaderStage()` returns `'vertex'`, `'fragment'` or `'compute'` when the function was transitively
   used in a vertex, fragment or compute entry function respectively, `undefined` otherwise.
+- `std.typeOf(value)` returns the schema of `value` as it will be typed in the generated shader (e.g. `d.f32`),
+  or `undefined` if it cannot be determined. It can only be called during shader generation, and its result has to be
+  consumed immediately (compared or passed as an argument), since it cannot be stored in a variable.
 
 ```ts twoslash
 import { tgpu, d, std } from 'typegpu';
diff --git a/packages/typegpu-testing-utility/src/index.ts b/packages/typegpu-testing-utility/src/index.ts
index 7d99f5da03..43117daa53 100644
--- a/packages/typegpu-testing-utility/src/index.ts
+++ b/packages/typegpu-testing-utility/src/index.ts
@@ -1,3 +1,35 @@
+import { expect } from 'vitest';
+
+declare module 'vitest' {
+  interface Matchers<T> {
+    toMatchNTimes: (regexp: RegExp, expectedMatches: number) => T;
+  }
+}
+
+expect.extend({
+  toMatchNTimes(received, regexp, expectedMatches) {
+    const { isNot } = this;
+
+    // A fresh global regex, so that non-global patterns don't loop forever,
+    // and the caller's `lastIndex` is left untouched.
+    const globalRegexp = new RegExp(
+      regexp.source,
+      regexp.flags.includes('g') ? regexp.flags : `${regexp.flags}g`,
+    );
+    const occurrences = [...String(received).matchAll(globalRegexp)].length;
+
+    return {
+      pass: occurrences === expectedMatches,
+      actual: occurrences,
+      expected: expectedMatches,
+      message: () =>
+        isNot
+          ? `${received} shouldn't contain the pattern ${expectedMatches} times`
+          : `${received} should contain the pattern ${expectedMatches} times`,
+    };
+  },
+});
+
 export {
   it,
   test,
diff --git a/packages/typegpu/src/data/index.ts b/packages/typegpu/src/data/index.ts
index 8bcad1cfd9..269ff371bf 100644
--- a/packages/typegpu/src/data/index.ts
+++ b/packages/typegpu/src/data/index.ts
@@ -46,7 +46,7 @@ export const mat2x2f = (() => (assignInfixOperators(), _mat2x2f))();
 export const mat3x3f = (() => (assignInfixOperators(), _mat3x3f))();
 export const mat4x4f = (() => (assignInfixOperators(), _mat4x4f))();
 
-export { bool, f16, f32, i32, u16, u32 } from './numeric.ts';
+export { abstractFloat, abstractInt, bool, f16, f32, i32, u16, u32 } from './numeric.ts';
 export {
   isAlignAttrib,
   isAtomic,
@@ -63,6 +63,8 @@ export {
   Void,
 } from './wgslTypes.ts';
 export type {
+  AbstractFloat,
+  AbstractInt,
   Align,
   AnyVecInstance,
   AnyWgslData,
diff --git a/packages/typegpu/src/std/index.ts b/packages/typegpu/src/std/index.ts
index b3c3e40cdf..da047930b6 100644
--- a/packages/typegpu/src/std/index.ts
+++ b/packages/typegpu/src/std/index.ts
@@ -215,3 +215,5 @@ export { bitcastU32toF32, bitcastU32toI32, bitcastF32toU32, bitcast } from './bi
 export { range } from './range.ts';
 
 export { isBeingTranspiled, getTargetShaderLanguage, getShaderStage } from './environment.ts';
+
+export { typeOf } from './typeOf.ts';
diff --git a/packages/typegpu/src/std/typeOf.ts b/packages/typegpu/src/std/typeOf.ts
new file mode 100644
index 0000000000..b506cee859
--- /dev/null
+++ b/packages/typegpu/src/std/typeOf.ts
@@ -0,0 +1,68 @@
+import type { AnyData } from '../data/dataTypes.ts';
+import { isWgslData } from '../data/wgslTypes.ts';
+import { setName } from '../shared/meta.ts';
+import { $gpuCallable, $internal } from '../shared/symbols.ts';
+import { coerceToSnippet } from '../tgsl/generationHelpers.ts';
+import type { GPUCallable } from '../types.ts';
+
+interface TypeOf extends GPUCallable {
+  (arg: unknown): AnyData | undefined;
+}
+
+/**
+ * Returns the type of an expression, the one it will have during shader execution. That information can
+ * be used to introspect or specialize the shader. It will return `undefined` if the expression cannot
+ * be represented in shader code, or if the type is unknown at the point where the query is made.
+ *
+ * The returned value is only known at shader-generation time and has no WGSL representation, so it
+ * cannot be stored in a variable (e.g. `const t = std.typeOf(a)` throws). Consume it immediately,
+ * e.g. by comparing it (`std.typeOf(a) === d.f32`) or passing it to a comptime function.
+ *
+ * @example
+ * function double(value: number) {
+ *   'use gpu';
+ *   if (std.typeOf(value) === d.f32) {
+ *     return value * 2;
+ *   }
+ *   return value << 1;
+ * }
+ *
+ * function main() {
+ *   'use gpu';
+ *   double(1.5); // generates double(f32), which multiplies by 2
+ *   double(1); // generates double(i32), which shifts left by 1
+ * }
+ */
+export const typeOf = ((): TypeOf => {
+  const impl = (() => {
+    // Deducing the WGSL type of a value at comptime is near-impossible with
+    // the current comptime type flowing. We can lift this restriction once it
+    // becomes more capable.
+    throw new Error(
+      'std.typeOf() cannot be called in JS, it can only be used in the GPU context. You can use std.isBeingTranspiled() to branch off and provide a fallback.',
+    );
+  }) as unknown as TypeOf;
+
+  setName(impl, 'typeOf');
+  impl.toString = () => 'typeOf';
+  impl[$gpuCallable] = {
+    call(_ctx, args) {
+      const [arg] = args;
+      if (!arg || args.length > 1) {
+        throw new Error(`std.typeOf() expects exactly one argument, got ${args.length}`);
+      }
+      const type = arg.dataType;
+      // Only public data types are exposed. Unknown types and internal types
+      // (e.g. AutoStruct, EntryInputRouter) must not leak to the user.
+      if (!isWgslData(type)) {
+        return coerceToSnippet(undefined);
+      }
+
+      return coerceToSnippet(type);
+    },
+  };
+  // Mark as internal
+  Object.defineProperty(impl, $internal, { value: true });
+
+  return impl;
+})();
diff --git a/packages/typegpu/tests/std/typeOf.test.ts b/packages/typegpu/tests/std/typeOf.test.ts
new file mode 100644
index 0000000000..b3dbbd124d
--- /dev/null
+++ b/packages/typegpu/tests/std/typeOf.test.ts
@@ -0,0 +1,354 @@
+import { test } from 'typegpu-testing-utility';
+import { d, std, tgpu } from 'typegpu';
+import { describe, expect } from 'vitest';
+
+test('std.typeOf of a scalar argument', () => {
+  function foo(a: number, b: number, c: boolean) {
+    'use gpu';
+    let result = true;
+    // All of these should be true
+    result = result || std.typeOf(a) === d.f32;
+    result = result || std.typeOf(b) === d.i32;
+    result = result || std.typeOf(c) === d.bool;
+    return result;
+  }
+
+  function main() {
+    'use gpu';
+    foo(1.5, 1, true);
+  }
+
+  const code = tgpu.resolve([main]);
+
+  expect(code).toMatchNTimes(/result = \(result \|\| true\)/g, 3);
+
+  expect(code).toMatchInlineSnapshot(`
+    "fn foo(a: f32, b: i32, c: bool) -> bool {
+      var result = true;
+      result = (result || true);
+      result = (result || true);
+      result = (result || true);
+      return result;
+    }
+
+    fn main() {
+      foo(1.5f, 1i, true);
+    }"
+  `);
+});
+
+test('std.typeOf to assert argument types', () => {
+  const assertType = tgpu.comptime((received: d.AnyData | undefined, expected: d.AnyData) => {
+    if (received !== expected) {
+      throw new Error(`Expected type ${String(expected)}, got ${String(received)}`);
+    }
+  });
+
+  function foo(a: number) {
+    'use gpu';
+    assertType(std.typeOf(a), d.f32);
+    return a * 2;
+  }
+
+  function good() {
+    'use gpu';
+    return foo(1.5);
+  }
+
+  function bad() {
+    'use gpu';
+    return foo(1);
+  }
+
+  expect(() => tgpu.resolve([good])).not.toThrow();
+  expect(() => tgpu.resolve([bad])).toThrowErrorMatchingInlineSnapshot(`
+    [Error: Resolution of the following tree failed:
+    - <root>
+    - fn*:bad
+    - fn*:bad()
+    - fn*:foo(i32)
+    - fn:assertType: Expected type f32, got i32]
+  `);
+});
+
+test('std.typeOf throws at comptime', () => {
+  expect(() => std.typeOf(true)).toThrowErrorMatchingInlineSnapshot(
+    `[Error: std.typeOf() cannot be called in JS, it can only be used in the GPU context. You can use std.isBeingTranspiled() to branch off and provide a fallback.]`,
+  );
+
+  const getType = tgpu.comptime((value: number) => std.typeOf(value));
+
+  function main() {
+    'use gpu';
+    getType(1.5);
+  }
+
+  expect(() => tgpu.resolve([main])).toThrowErrorMatchingInlineSnapshot(`
+    [Error: Resolution of the following tree failed:
+    - <root>
+    - fn*:main
+    - fn*:main()
+    - fn:getType: std.typeOf() cannot be called in JS, it can only be used in the GPU context. You can use std.isBeingTranspiled() to branch off and provide a fallback.]
+  `);
+});
+
+test('std.typeOf of a constant', () => {
+  const fn = () => {
+    'use gpu';
+    return std.typeOf(0) === d.abstractInt;
+  };
+
+  const code = tgpu.resolve([fn]);
+  expect(code).toMatchInlineSnapshot(`
+    "fn fn_1() -> bool {
+      return true;
+    }"
+  `);
+  expect(code).toContain('true');
+});
+
+test('std.typeOf of a struct', () => {
+  const Schema = d.struct({ prop: d.u32 });
+  const fn = () => {
+    'use gpu';
+    return std.typeOf(Schema()) === Schema;
+  };
+
+  const code = tgpu.resolve([fn]);
+  expect(code).toMatchInlineSnapshot(`
+    "struct Schema {
+      prop: u32,
+    }
+
+    fn fn_1() -> bool {
+      return true;
+    }"
+  `);
+  expect(code).toContain('true');
+});
+
+test('std.typeOf of UnknownData', () => {
+  const fn = () => {
+    'use gpu';
+    const t1 = std.typeOf('unknown') === undefined;
+    const t2 = std.typeOf(fn) === undefined;
+    const t3 = std.typeOf(d.vec2f) === undefined;
+    const t4 = std.typeOf(d.vec2f().mul) === undefined;
+    return t1 && t2 && t3 && t4;
+  };
+
+  const code = tgpu.resolve([fn]);
+  expect(code).toMatchInlineSnapshot(`
+    "fn fn_1() -> bool {
+      const t1 = true;
+      const t2 = true;
+      const t3 = true;
+      const t4 = true;
+      return (((t1 && t2) && t3) && t4);
+    }"
+  `);
+  expect(code).not.toContain('false');
+});
+
+test('std.typeOf assigned to a variable throws', () => {
+  const fn = () => {
+    'use gpu';
+    const t = std.typeOf(0);
+    return t;
+  };
+
+  expect(() => tgpu.resolve([fn])).toThrowErrorMatchingInlineSnapshot(`
+    [Error: Resolution of the following tree failed:
+    - <root>
+    - fn*:fn
+    - fn*:fn(): 'const t = std.typeOf(0)' is invalid, cannot determine WGSL type of 'std.typeOf(0)'
+    -----
+    - Try using or defining a schema that matches your desired value the most, and wrap the value with it: 'const t = Schema(std.typeOf(0))'
+    -----]
+  `);
+});
+
+test('std.typeOf of an AutoStruct is undefined', ({ root }) => {
+  const types: unknown[] = [];
+  const capture = tgpu.comptime((type: d.AnyData | undefined) => {
+    types.push(type);
+  });
+
+  const pipeline = root.createRenderPipeline({
+    vertex: () => {
+      'use gpu';
+      return { $position: d.vec4f(), uv: d.vec2f() };
+    },
+    fragment: (input) => {
+      'use gpu';
+      capture(std.typeOf(input));
+      return d.vec4f(input.uv, 0, 1);
+    },
+  });
+
+  tgpu.resolve([pipeline]);
+  expect(types).toStrictEqual([undefined]);
+});
+
+test('std.typeOf of an entry function input is undefined', () => {
+  const types: unknown[] = [];
+  const capture = tgpu.comptime((type: d.AnyData | undefined) => {
+    types.push(type);
+  });
+
+  const fragment = tgpu.fragmentFn({
+    in: { uv: d.vec2f },
+    out: d.vec4f,
+  })((input) => {
+    'use gpu';
+    capture(std.typeOf(input));
+    return d.vec4f(input.uv, 0, 1);
+  });
+
+  tgpu.resolve([fragment]);
+  expect(types).toStrictEqual([undefined]);
+});
+
+describe('std.typeOf of an abstruct', () => {
+  test('returns the result struct of std.frexp', () => {
+    const types: unknown[] = [];
+    const capture = tgpu.comptime((type: d.AnyData | undefined) => {
+      types.push(type);
+    });
+
+    const fn = (value: number) => {
+      'use gpu';
+      const result = std.frexp(value);
+      capture(std.typeOf(result));
+      capture(std.typeOf(std.frexp(d.vec3f(value))));
+      capture(std.typeOf(std.frexp(1.5)));
+    };
+
+    const main = () => {
+      'use gpu';
+      fn(1.5);
+    };
+
+    tgpu.resolve([main]);
+
+    const [f32Result, vec3fResult, abstractResult] = types as d.AnyWgslStruct[];
+    expect(f32Result?.type).toBe('struct');
+    expect(f32Result?.propTypes).toStrictEqual({ fract: d.f32, exp: d.i32 });
+    expect(vec3fResult?.type).toBe('struct');
+    expect(vec3fResult?.propTypes).toStrictEqual({ fract: d.vec3f, exp: d.vec3i });
+    expect(abstractResult?.type).toBe('struct');
+    expect(abstractResult?.propTypes).toStrictEqual({
+      fract: d.abstractFloat,
+      exp: d.abstractInt,
+    });
+  });
+
+  test('returns the result struct of std.modf', () => {
+    const types: unknown[] = [];
+    const capture = tgpu.comptime((type: d.AnyData | undefined) => {
+      types.push(type);
+    });
+
+    const fn = () => {
+      'use gpu';
+      capture(std.typeOf(std.modf(d.f32(1.5))));
+      capture(std.typeOf(std.modf(d.vec2h(1.5))));
+    };
+
+    tgpu.resolve([fn]);
+
+    const [f32Result, vec2hResult] = types as d.AnyWgslStruct[];
+    expect(f32Result?.propTypes).toStrictEqual({ fract: d.f32, whole: d.f32 });
+    expect(vec2hResult?.propTypes).toStrictEqual({ fract: d.vec2h, whole: d.vec2h });
+  });
+
+  test('returns the result struct of std.atomicCompareExchangeWeak', ({ root }) => {
+    const types: unknown[] = [];
+    const capture = tgpu.comptime((type: d.AnyData | undefined) => {
+      types.push(type);
+    });
+
+    const counter = root.createMutable(d.atomic(d.u32));
+
+    const fn = () => {
+      'use gpu';
+      capture(std.typeOf(std.atomicCompareExchangeWeak(counter.$, 0, 1)));
+    };
+
+    tgpu.resolve([fn]);
+
+    const [result] = types as d.AnyWgslStruct[];
+    expect(result?.type).toBe('struct');
+    expect(result?.propTypes).toStrictEqual({ old_value: d.u32, exchanged: d.bool });
+  });
+
+  test('can be compared at comptime', () => {
+    const fn = (a: number, b: number) => {
+      'use gpu';
+      const same = std.typeOf(std.frexp(a)) === std.typeOf(std.frexp(b));
+      const different = std.typeOf(std.frexp(a)) === std.typeOf(std.modf(a));
+      return same && !different;
+    };
+
+    const main = () => {
+      'use gpu';
+      fn(1.5, 2.5);
+    };
+
+    expect(tgpu.resolve([main])).toMatchInlineSnapshot(`
+      "fn fn_1(a: f32, b: f32) -> bool {
+        const same = true;
+        const different = false;
+        return (same && !(different));
+      }
+
+      fn main() {
+        fn_1(1.5f, 2.5f);
+      }"
+    `);
+  });
+
+  test('of a field access returns the field type', () => {
+    const fn = (value: number) => {
+      'use gpu';
+      const result = std.frexp(value);
+      const t1 = std.typeOf(result.fract) === d.f32;
+      const t2 = std.typeOf(result.exp) === d.i32;
+      return t1 && t2;
+    };
+
+    const main = () => {
+      'use gpu';
+      fn(1.5);
+    };
+
+    const code = tgpu.resolve([main]);
+    expect(code).toMatchInlineSnapshot(`
+      "fn fn_1(value: f32) -> bool {
+        let result = frexp(value);
+        const t1 = true;
+        const t2 = true;
+        return (t1 && t2);
+      }
+
+      fn main() {
+        fn_1(1.5f);
+      }"
+    `);
+    expect(code).not.toContain('false');
+  });
+
+  test('does not emit a struct declaration', () => {
+    const fn = (value: number) => {
+      'use gpu';
+      return std.typeOf(std.frexp(value)) === std.typeOf(std.frexp(value));
+    };
+
+    const main = () => {
+      'use gpu';
+      fn(1.5);
+    };
+
+    expect(tgpu.resolve([main])).not.toContain('struct');
+  });
+});