feat: std.typeOf

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';

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,

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,

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';

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