chore: Add type checks to std

packages/typegpu/src/data/generalizeFn.ts

@@ -12,7 +12,7 @@ import {
   type v3b,
   type v4b,
 } from './wgslTypes.ts';
-import { invariant } from '../errors.ts';
+import { invariant, WgslTypeError } from '../errors.ts';
 
 type Vec = AnyVecInstance; // alias
 type Mat = AnyMatInstance; // alias
@@ -117,6 +117,10 @@ export function generalizeFn<T extends Algebraic>(fn: (...args: number[]) => num
 /**
  * Analogous to `generalizeFn`, but the return type is a boolean vector instead.
  */
+export function generalizeBoolFn<T extends boolean | AnyBooleanVecInstance>(
+  fn: (a: boolean) => boolean,
+  args: [T],
+): T;
 export function generalizeBoolFn<T extends Algebraic>(
   fn: (a: number, b: number) => boolean,
   args: [T, T],
@@ -132,7 +136,7 @@ export function generalizeBoolFn<T extends Algebraic>(
   return applyArgs(fn, args, 'boolean') as ToBool<T>;
 }
 
-function kindOf(v: Algebraic): Kind {
+export function kindOf(v: Algebraic): Kind {
   if (typeof v === 'number') {
     return 'number';
   }
@@ -142,11 +146,61 @@ function kindOf(v: Algebraic): Kind {
   return v.kind;
 }
 
+// Unless matrix is mentioned in the name, it is not included.
+const i32Kind: Set<Kind> = new Set(['number', 'vec2i', 'vec3i', 'vec4i']);
+export const u32Kind: Set<Kind> = new Set(['number', 'vec2u', 'vec3u', 'vec4u']);
+export const f32Kind: Set<Kind> = new Set(['number', 'vec2f', 'vec3f', 'vec4f']);
+export const f16Kind: Set<Kind> = new Set(['number', 'vec2h', 'vec3h', 'vec4h']);
+export const matrixKind: Set<Kind> = new Set(['mat2x2f', 'mat3x3f', 'mat4x4f']);
+export const booleanKind: Set<Kind> = new Set([
+  'boolean',
+  'vec2<bool>',
+  'vec3<bool>',
+  'vec4<bool>',
+]);
+export const floatKind: Set<Kind> = new Set([...f32Kind, ...f16Kind]);
+export const signedKind: Set<Kind> = new Set([...i32Kind, ...f32Kind, ...f16Kind]);
+export const numericKind: Set<Kind> = new Set([...signedKind, ...u32Kind]);
+export const numericOrBooleanKind: Set<Kind> = new Set([...numericKind, ...booleanKind]);
+export const numericOrMatrixKind: Set<Kind> = new Set([...numericKind, ...matrixKind]);
+export const crossKind: Set<Kind> = new Set(['vec3f', 'vec3h']);
+
+export function assertKind(
+  v: Algebraic | Algebraic[],
+  valid: Set<Kind>,
+  excludeScalar: boolean = false,
+) {
+  if (!isVecInstance(v) && Array.isArray(v)) {
+    v.forEach((item) => assertKind(item, valid, excludeScalar));
+    return;
+  }
+  const kind = kindOf(v);
+  if (!valid.has(kind)) {
+    throw new WgslTypeError(
+      `Unsupported signature. Expected one of '${[...valid].join(', ')}', got '${kind}'.`,
+    );
+  }
+  if (excludeScalar && (kind === 'number' || kind === 'boolean')) {
+    throw new WgslTypeError(
+      `Unsupported signature. Expected kind to not be scalar, got '${kind}'.`,
+    );
+  }
+}
+
+export function assertEqualKinds(...values: Algebraic[]) {
+  const kinds = new Set(values.map(kindOf));
+  if (kinds.size !== 1) {
+    throw new WgslTypeError(
+      `Unsupported signature. Expected the following kinds to be equal: '${[...kinds].join(', ')}'.`,
+    );
+  }
+}
+
 /**
  * If one of the arguments is a vector and other is a number,
  * the number is up-cased to a vector.
  */
-export function upCast<T extends number | Vec>(
+export function upCast<T extends number | Vec | Mat>(
   args: [T, T],
 ): [Exclude<T, number>, Exclude<T, number>] {
   const [lhs, rhs] = args;

packages/typegpu/src/std/bitcast.ts

@@ -44,6 +44,7 @@ import { getName } from '../shared/meta.ts';
 import type { Infer } from '../shared/repr.ts';
 import { comptime } from '../core/function/comptime.ts';
 import { coerceToSnippet } from '../tgsl/generationHelpers.ts';
+import { f32Kind, u32Kind, assertKind } from '../data/generalizeFn.ts';
 
 type BitcastU32toF32Overload = <T extends number | v2u | v3u | v4u>(
   value: T,
@@ -59,6 +60,7 @@ const u32AllowedSchemas = [u32, vec2u, vec3u, vec4u];
 export const bitcastU32toF32 = dualImpl({
   name: 'bitcastU32toF32',
   normalImpl: ((value) => {
+    assertKind(value, u32Kind);
     if (typeof value === 'number') {
       return bitcastU32toF32Impl(value);
     }
@@ -96,6 +98,7 @@ type BitcastU32toI32Overload = <T extends number | v2u | v3u | v4u>(
 export const bitcastU32toI32 = dualImpl({
   name: 'bitcastU32toI32',
   normalImpl: ((value) => {
+    assertKind(value, u32Kind);
     if (typeof value === 'number') {
       return bitcastU32toI32Impl(value);
     }
@@ -135,6 +138,7 @@ const f32AllowedSchemas = [f32, vec2f, vec3f, vec4f];
 export const bitcastF32toU32 = dualImpl({
   name: 'bitcastF32toU32',
   normalImpl: ((value) => {
+    assertKind(value, f32Kind);
     if (typeof value === 'number') {
       return bitcastF32toU32Impl(value);
     }

packages/typegpu/src/std/boolean.ts

@@ -20,7 +20,17 @@ import {
   vec4u,
 } from '../data/vector.ts';
 import { VectorOps } from '../data/vectorOps.ts';
-import { generalizeBoolFn, generalizeFn } from '../data/generalizeFn.ts';
+import {
+  booleanKind,
+  floatKind,
+  generalizeBoolFn,
+  generalizeFn,
+  kindOf,
+  numericKind,
+  numericOrBooleanKind,
+  assertEqualKinds,
+  assertKind,
+} from '../data/generalizeFn.ts';
 import {
   type AnyBooleanVecInstance,
   type AnyFloatVecInstance,
@@ -32,12 +42,12 @@ import {
   type BaseData,
   isBool,
   isVecBool,
-  isVecBoolInstance,
+  isVecInstance,
   type v2b,
   type v3b,
   type v4b,
 } from '../data/wgslTypes.ts';
-import { SignatureNotSupportedError } from '../errors.ts';
+import { SignatureNotSupportedError, WgslTypeError } from '../errors.ts';
 import { unify } from '../tgsl/conversion.ts';
 import { cpuCopy } from './copy.ts';
 
@@ -68,8 +78,11 @@ export const allEq = dualImpl({
   sideEffects: false,
 });
 
-const cpuEq = <T extends AnyVecInstance>(lhs: T, rhs: T) =>
-  generalizeBoolFn((a, b) => a === b, [lhs, rhs]);
+const cpuEq = <T extends AnyVecInstance>(lhs: T, rhs: T) => {
+  assertKind([lhs, rhs], numericOrBooleanKind);
+  assertEqualKinds(lhs, rhs);
+  return generalizeBoolFn((a, b) => a === b, [lhs, rhs]);
+};
 
 /**
  * Checks **component-wise** whether `lhs == rhs`.
@@ -110,8 +123,11 @@ export const ne = dualImpl({
   sideEffects: false,
 });
 
-const cpuLt = <T extends AnyNumericVecInstance>(lhs: T, rhs: T) =>
-  generalizeBoolFn((a, b) => a < b, [lhs, rhs]);
+const cpuLt = <T extends AnyNumericVecInstance>(lhs: T, rhs: T) => {
+  assertKind([lhs, rhs], numericKind);
+  assertEqualKinds(lhs, rhs);
+  return generalizeBoolFn((a, b) => a < b, [lhs, rhs]);
+};
 
 /**
  * Checks **component-wise** whether `lhs < rhs`.
@@ -196,22 +212,8 @@ export const ge = dualImpl({
 function cpuNot(value: boolean): boolean;
 function cpuNot<T extends AnyBooleanVecInstance>(value: T): T;
 function cpuNot<T extends AnyBooleanVecInstance | boolean>(value: T): T {
-  if (typeof value === 'boolean') {
-    return !value as T;
-  }
-
-  if (!isVecBoolInstance(value)) {
-    throw new Error(`'std.not' requires a boolean or boolean vector.`);
-  }
-
-  switch (value.length) {
-    case 2:
-      return vec2b(cpuNot(value.x), cpuNot(value.y)) as T;
-    case 3:
-      return vec3b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z)) as T;
-    case 4:
-      return vec4b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z), cpuNot(value.w)) as T;
-  }
+  assertKind(value, booleanKind);
+  return generalizeBoolFn((a: boolean) => !a, [value]);
 }
 
 /**
@@ -239,8 +241,11 @@ export const not = dualImpl({
   sideEffects: false,
 });
 
-const cpuOr = <T extends AnyBooleanVecInstance>(lhs: T, rhs: T) =>
-  generalizeBoolFn((a: boolean, b: boolean) => a || b, [lhs, rhs]);
+const cpuOr = <T extends AnyBooleanVecInstance>(lhs: T, rhs: T) => {
+  assertKind([lhs, rhs], booleanKind);
+  assertEqualKinds(lhs, rhs);
+  return generalizeBoolFn((a: boolean, b: boolean) => a || b, [lhs, rhs]);
+};
 
 /**
  * Returns **component-wise** logical `or` result.
@@ -275,7 +280,13 @@ export const and = dualImpl({
 
 // logical aggregation
 
-const cpuAll = (value: AnyBooleanVecInstance) => VectorOps.all[value.kind](value);
+const cpuAll = (value: boolean | AnyBooleanVecInstance) => {
+  assertKind(value, booleanKind);
+  if (typeof value === 'boolean') {
+    return value;
+  }
+  return VectorOps.all[value.kind](value);
+};
 
 /**
  * Returns `true` if each component of `value` is true.
@@ -328,6 +339,8 @@ export const isCloseTo = dualImpl({
     rhs: T,
     precision = 0.01,
   ): boolean => {
+    assertKind([lhs, rhs], floatKind);
+    assertEqualKinds(lhs, rhs);
     const componentResult = generalizeBoolFn(
       (lhs, rhs) => Math.abs(lhs - rhs) < precision,
       [lhs, rhs],
@@ -361,9 +374,17 @@ function cpuSelect<T extends number | boolean | AnyVecInstance>(
   t: T,
   cond: AnyBooleanVecInstance | boolean,
 ) {
+  assertKind([f, t], numericOrBooleanKind);
+  assertEqualKinds(f, t);
+  assertKind(cond, booleanKind);
   if (typeof cond === 'boolean') {
     return cpuCopy(cond ? t : f);
   }
+  if (!isVecInstance(f) || f.length !== cond.length) {
+    throw new WgslTypeError(
+      `Select shape '(${kindOf(f)}, ${kindOf(t)}, ${kindOf(cond)})' is invalid.`,
+    );
+  }
   // generalizeFn will handle this fine, it just has no mixed type overload.
   return generalizeFn((f, t, c) => (c ? t : f), [f, t, cond as T]);
 }