feat: std.typeOf

packages/typegpu-testing-utility/src/index.ts

@@ -1 +1,31 @@
+import { expect } from 'vitest';
+
+declare module 'vitest' {
+  interface Matchers<T = any> {
+    toMatchNTimes: (regexp: RegExp, expectedMatches: number) => T;
+  }
+}
+
+expect.extend({
+  toMatchNTimes(received, regexp, expectedMatches) {
+    const { isNot } = this;
+
+    let occurrances = 0;
+
+    while (regexp.exec(received) !== null) {
+      occurrances++;
+    }
+
+    return {
+      pass: occurrances === expectedMatches,
+      actual: occurrances,
+      expected: expectedMatches,
+      message: () =>
+        isNot
+          ? `${received} shouldn't contain the pattern ${expectedMatches} times`
+          : `${received} should contain the pattern ${expectedMatches} times`,
+    };
+  },
+});
+
 export { it, test } from './extendedIt.ts';

packages/typegpu/src/std/index.ts

@@ -189,3 +189,5 @@ export { extensionEnabled } from './extensions.ts';
 export { bitcastU32toF32, bitcastU32toI32 } from './bitcast.ts';
 
 export { range } from './range.ts';
+
+export { typeOf } from './typeOf.ts';

packages/typegpu/src/std/typeOf.ts

@@ -0,0 +1,41 @@
+import { AutoStruct } from '../data/autoStruct.ts';
+import { UnknownData, type AnyData } from '../data/dataTypes.ts';
+import { bool } from '../data/numeric.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;
+}
+
+// TODO: Idea, maybe concretize types automatically?
+export const typeOf = ((): TypeOf => {
+  const impl: TypeOf = (arg) => {
+    // TODO: Determine more types from values
+    if (typeof arg === 'boolean') {
+      return bool;
+    }
+    return undefined;
+  };
+
+  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;
+      if (type === UnknownData || type instanceof AutoStruct) {
+        return undefined;
+      }
+
+      return coerceToSnippet(type);
+    },
+  };
+  // Mark as internal
+  Object.defineProperty(impl, $internal, {});
+
+  return impl;
+})();

packages/typegpu/tests/std/typeOf.test.ts

@@ -0,0 +1,72 @@
+import { test } from 'typegpu-testing-utility';
+import tgpu, { d, std } from 'typegpu';
+import { 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]
+  `);
+});