feat: bitcastF32toU32

packages/typegpu/src/data/numberOps.ts

@@ -16,14 +16,21 @@ export const divInteger = (lhs: number, rhs: number) => {
   return Math.trunc(lhs / rhs);
 };
 
+const buf32 = new ArrayBuffer(4);
+const f32arr = new Float32Array(buf32);
+const u32arr = new Uint32Array(buf32);
+const i32arr = new Int32Array(buf32);
 export function bitcastU32toF32Impl(n: number): number {
-  const dataView = new DataView(new ArrayBuffer(4));
-  dataView.setUint32(0, n, true);
-  return dataView.getFloat32(0, true);
+  u32arr[0] = n;
+  return f32arr[0] as number;
 }
 
 export function bitcastU32toI32Impl(n: number): number {
-  const dataView = new DataView(new ArrayBuffer(4));
-  dataView.setUint32(0, n, true);
-  return dataView.getInt32(0, true);
+  u32arr[0] = n;
+  return i32arr[0] as number;
+}
+
+export function bitcastF32toU32Impl(n: number): number {
+  f32arr[0] = n;
+  return u32arr[0] as number;
 }

packages/typegpu/src/data/vectorOps.ts

@@ -1,5 +1,6 @@
 import { mat2x2f, mat3x3f, mat4x4f } from './matrix.ts';
 import {
+  bitcastF32toU32Impl,
   bitcastU32toF32Impl,
   bitcastU32toI32Impl,
   clamp,
@@ -1162,4 +1163,22 @@ export const VectorOps = {
       v: T,
     ) => T extends wgsl.v2u ? wgsl.v2i : T extends wgsl.v3u ? wgsl.v3i : wgsl.v4i
   >,
+
+  bitcastF32toU32: {
+    vec2f: (n: wgsl.v2f) => vec2u(bitcastF32toU32Impl(n.x), bitcastF32toU32Impl(n.y)),
+    vec3f: (n: wgsl.v3f) =>
+      vec3u(bitcastF32toU32Impl(n.x), bitcastF32toU32Impl(n.y), bitcastF32toU32Impl(n.z)),
+    vec4f: (n: wgsl.v4f) =>
+      vec4u(
+        bitcastF32toU32Impl(n.x),
+        bitcastF32toU32Impl(n.y),
+        bitcastF32toU32Impl(n.z),
+        bitcastF32toU32Impl(n.w),
+      ),
+  } as Record<
+    VecKind,
+    <T extends wgsl.AnyFloatVecInstance>(
+      v: T,
+    ) => T extends wgsl.v2f ? wgsl.v2u : T extends wgsl.v3f ? wgsl.v3u : wgsl.v4u
+  >,
 };

packages/typegpu/src/std/bitcast.ts

@@ -1,17 +1,23 @@
 import { dualImpl } from '../core/function/dualImpl.ts';
 import { stitch } from '../core/resolve/stitch.ts';
-import { bitcastU32toF32Impl, bitcastU32toI32Impl } from '../data/numberOps.ts';
+import {
+  bitcastF32toU32Impl,
+  bitcastU32toF32Impl,
+  bitcastU32toI32Impl,
+} from '../data/numberOps.ts';
 import { f32, i32, u32 } from '../data/numeric.ts';
 import { isVec } from '../data/wgslTypes.ts';
-import { vec2f, vec2i, vec3f, vec3i, vec4f, vec4i } from '../data/vector.ts';
+import { vec2f, vec2i, vec2u, vec3f, vec3i, vec3u, vec4f, vec4i, vec4u } from '../data/vector.ts';
 import { VectorOps } from '../data/vectorOps.ts';
 import type { v2f, v2i, v2u, v3f, v3i, v3u, v4f, v4i, v4u } from '../data/wgslTypes.ts';
-import { unify } from '../tgsl/conversion.ts';
+import { unifyStrict } from '../tgsl/conversion.ts';
+import { SignatureNotSupportedError } from '../errors.ts';
 
-export type BitcastU32toF32Overload = ((value: number) => number) &
-  ((value: v2u) => v2f) &
-  ((value: v3u) => v3f) &
-  ((value: v4u) => v4f);
+type BitcastU32toF32Overload = <T extends number | v2u | v3u | v4u>(
+  value: T,
+) => T extends v2u ? v2f : T extends v3u ? v3f : T extends v4u ? v4f : number;
+
+const u32AllowedSchemas = [u32, vec2u, vec3u, vec4u];
 
 export const bitcastU32toF32 = dualImpl({
   name: 'bitcastU32toF32',
@@ -21,9 +27,16 @@ export const bitcastU32toF32 = dualImpl({
     }
     return VectorOps.bitcastU32toF32[value.kind](value);
   }) as BitcastU32toF32Overload,
-  codegenImpl: (_ctx, [n]) => stitch`bitcast<f32>(${n})`,
+  codegenImpl: (_ctx, [n]) => {
+    return isVec(n.dataType)
+      ? stitch`bitcast<vec${n.dataType.componentCount}f>(${n})`
+      : stitch`bitcast<f32>(${n})`;
+  },
   signature: (...arg) => {
-    const uargs = unify(arg, [u32]) ?? arg;
+    const uargs = unifyStrict(arg, u32AllowedSchemas);
+    if (!uargs) {
+      throw new SignatureNotSupportedError(arg, u32AllowedSchemas);
+    }
     return {
       argTypes: uargs,
       returnType: isVec(uargs[0])
@@ -37,10 +50,9 @@ export const bitcastU32toF32 = dualImpl({
   },
 });
 
-export type BitcastU32toI32Overload = ((value: number) => number) &
-  ((value: v2u) => v2i) &
-  ((value: v3u) => v3i) &
-  ((value: v4u) => v4i);
+type BitcastU32toI32Overload = <T extends number | v2u | v3u | v4u>(
+  value: T,
+) => T extends v2u ? v2i : T extends v3u ? v3i : T extends v4u ? v4i : number;
 
 export const bitcastU32toI32 = dualImpl({
   name: 'bitcastU32toI32',
@@ -50,9 +62,16 @@ export const bitcastU32toI32 = dualImpl({
     }
     return VectorOps.bitcastU32toI32[value.kind](value);
   }) as BitcastU32toI32Overload,
-  codegenImpl: (_ctx, [n]) => stitch`bitcast<i32>(${n})`,
+  codegenImpl: (_ctx, [n]) => {
+    return isVec(n.dataType)
+      ? stitch`bitcast<vec${n.dataType.componentCount}i>(${n})`
+      : stitch`bitcast<i32>(${n})`;
+  },
   signature: (...arg) => {
-    const uargs = unify(arg, [u32]) ?? arg;
+    const uargs = unifyStrict(arg, u32AllowedSchemas);
+    if (!uargs) {
+      throw new SignatureNotSupportedError(arg, u32AllowedSchemas);
+    }
     return {
       argTypes: uargs,
       returnType: isVec(uargs[0])
@@ -65,3 +84,40 @@ export const bitcastU32toI32 = dualImpl({
     };
   },
 });
+
+type BitcastF32toU32Overload = <T extends number | v2f | v3f | v4f>(
+  value: T,
+) => T extends v2f ? v2u : T extends v3f ? v3u : T extends v4f ? v4u : number;
+
+const f32AllowedSchemas = [f32, vec2f, vec3f, vec4f];
+
+export const bitcastF32toU32 = dualImpl({
+  name: 'bitcastF32toU32',
+  normalImpl: ((value) => {
+    if (typeof value === 'number') {
+      return bitcastF32toU32Impl(value);
+    }
+    return VectorOps.bitcastF32toU32[value.kind](value);
+  }) as BitcastF32toU32Overload,
+  codegenImpl: (_ctx, [n]) => {
+    return isVec(n.dataType)
+      ? stitch`bitcast<vec${n.dataType.componentCount}u>(${n})`
+      : stitch`bitcast<u32>(${n})`;
+  },
+  signature: (...arg) => {
+    const uargs = unifyStrict(arg, f32AllowedSchemas);
+    if (!uargs) {
+      throw new SignatureNotSupportedError(arg, f32AllowedSchemas);
+    }
+    return {
+      argTypes: uargs,
+      returnType: isVec(uargs[0])
+        ? uargs[0].type === 'vec2f'
+          ? vec2u
+          : uargs[0].type === 'vec3f'
+            ? vec3u
+            : vec4u
+        : u32,
+    };
+  },
+});

packages/typegpu/src/std/index.ts

@@ -186,7 +186,7 @@ export {
 
 export { extensionEnabled } from './extensions.ts';
 
-export { bitcastU32toF32, bitcastU32toI32 } from './bitcast.ts';
+export { bitcastU32toF32, bitcastU32toI32, bitcastF32toU32 } from './bitcast.ts';
 
 export { range } from './range.ts';
 

packages/typegpu/src/tgsl/conversion.ts

@@ -321,6 +321,25 @@ export function unify<T extends (BaseData | UnknownData)[] | []>(
   };
 }
 
+export function unifyStrict<T extends (BaseData | UnknownData)[] | []>(
+  inTypes: T,
+  restrictTo?: BaseData[],
+): { [K in keyof T]: BaseData } | undefined {
+  if (inTypes.some((type) => type === UnknownData)) {
+    return undefined;
+  }
+
+  const uniqueTargetTypes = [...new Set(((restrictTo || inTypes) as BaseData[]).map(undecorate))];
+  const conversion = findBestType(inTypes as BaseData[], uniqueTargetTypes, false);
+  if (!conversion) {
+    return undefined;
+  }
+
+  return inTypes.map((type) => (isVec(type) || isMat(type) ? type : conversion.targetType)) as {
+    [K in keyof T]: BaseData;
+  };
+}
+
 export function convertToCommonType<T extends Snippet[]>(
   ctx: ResolutionCtx,
   values: T,