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',
@@ -23,7 +29,10 @@ export const bitcastU32toF32 = dualImpl({
   }) as BitcastU32toF32Overload,
   codegenImpl: (_ctx, [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 +46,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',
@@ -52,7 +60,10 @@ export const bitcastU32toI32 = dualImpl({
   }) as BitcastU32toI32Overload,
   codegenImpl: (_ctx, [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 +76,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,

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

@@ -12,6 +12,9 @@ import {
 } from '../../src/data/vector.ts';
 import tgpu, { d, std } from '../../src/index.js';
 
+// remember to pad with zeros to 8 hex symbols
+const floatFromHex = (hex: string) => Buffer.from(hex, 'hex').readFloatBE(0);
+
 describe('bitcast', () => {
   it('bitcastU32toF32', () => {
     // 1.0 in f32
@@ -37,6 +40,14 @@ describe('bitcast', () => {
     expect(i2).toBe(-2147483648);
   });
 
+  it('bitcastF32toU32', () => {
+    const i1 = std.bitcastF32toU32(floatFromHex('00000001'));
+    expect(i1).toBe(1);
+
+    const i2 = std.bitcastF32toU32(floatFromHex('7f800000'));
+    expect(i2).toBe(2139095040);
+  });
+
   it('bitcastU32toF32 vectors', () => {
     const v2 = vec2u(1065353216, 3212836864); // 1.0f, -1.0f
     const cast2 = std.bitcastU32toF32(v2);
@@ -65,6 +76,25 @@ describe('bitcast', () => {
     expect(cast4).toEqual(vec4i(0, 1, -1, -2147483648));
   });
 
+  it('bitcastF32toU32 vectors', () => {
+    const v2 = vec2f(floatFromHex('7f800000'), floatFromHex('7fc00000')); // +inf, quiet nan
+    const cast2 = std.bitcastF32toU32(v2);
+    expect(cast2).toStrictEqual(vec2u(2139095040, 2143289344));
+
+    const v3 = vec3f(floatFromHex('ff800000'), floatFromHex('00000001'), floatFromHex('80000001'));
+    const cast3 = std.bitcastF32toU32(v3);
+    expect(cast3).toStrictEqual(vec3u(4286578688, 1, 2147483649));
+
+    const v4 = vec4f(
+      floatFromHex('84220925'),
+      floatFromHex('68800000'),
+      floatFromHex('48980780'),
+      floatFromHex('0000075a'),
+    );
+    const cast4 = std.bitcastF32toU32(v4);
+    expect(cast4).toStrictEqual(vec4u(2216823077, 1753219072, 1217922944, 1882));
+  });
+
   it('bitcastU32toF32 specials (NaN, infinities etc)', () => {
     // +0
     const pz = std.bitcastU32toF32(0x00000000);
@@ -119,6 +149,7 @@ describe('bitcast in shaders', () => {
   it('works for primitives', () => {
     const fnf32 = tgpu.fn([], d.f32)(() => std.bitcastU32toF32(1234));
     const fni32 = tgpu.fn([], d.i32)(() => std.bitcastU32toI32(d.u32(2 ** 31)));
+    const fnu32 = tgpu.fn([d.f32], d.u32)((v) => std.bitcastF32toU32(v));
 
     expect(tgpu.resolve([fnf32])).toMatchInlineSnapshot(`
       "fn fnf32() -> f32 {
@@ -130,11 +161,17 @@ describe('bitcast in shaders', () => {
         return -2147483648i;
       }"
     `);
+    expect(tgpu.resolve([fnu32])).toMatchInlineSnapshot(`
+      "fn fnu32(v: f32) -> u32 {
+        return bitcast<u32>(v);
+      }"
+    `);
   });
 
   it('works for vectors', () => {
     const fnvec4i = tgpu.fn([], d.vec4i)(() => std.bitcastU32toI32(vec4u(1, 2, 3, 4)));
     const fnvec4f = tgpu.fn([], d.vec4f)(() => std.bitcastU32toF32(vec4u(1, 2, 3, 4)));
+    const fnvec4u = tgpu.fn([d.vec4f], d.vec4u)((v) => std.bitcastF32toU32(v));
 
     expect(tgpu.resolve([fnvec4i])).toMatchInlineSnapshot(`
       "fn fnvec4i() -> vec4i {
@@ -146,5 +183,64 @@ describe('bitcast in shaders', () => {
         return vec4f(1.401298464324817e-45, 2.802596928649634e-45, 4.203895392974451e-45, 5.605193857299268e-45);
       }"
     `);
+    expect(tgpu.resolve([fnvec4u])).toMatchInlineSnapshot(`
+      "fn fnvec4u(v: vec4f) -> vec4u {
+        return bitcast<vec4u>(v);
+      }"
+    `);
+  });
+
+  it('throws an error for unsupported signatures', () => {
+    const f1 = () => {
+      'use gpu';
+      // @ts-expect-error
+      return std.bitcastU32toF32(d.vec2i());
+    };
+    expect(() => tgpu.resolve([f1])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:f1
+      - fn*:f1()
+      - fn:bitcastU32toF32: Unsupported data types: vec2i. Supported types are: u32, vec2u, vec3u, vec4u.]
+    `);
+
+    const f2 = () => {
+      'use gpu';
+      // @ts-expect-error
+      return std.bitcastU32toI32(d.vec3f());
+    };
+    expect(() => tgpu.resolve([f2])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:f2
+      - fn*:f2()
+      - fn:bitcastU32toI32: Unsupported data types: vec3f. Supported types are: u32, vec2u, vec3u, vec4u.]
+    `);
+
+    const f3 = () => {
+      'use gpu';
+      // @ts-expect-error
+      return std.bitcastF32toU32(d.vec2h());
+    };
+    expect(() => tgpu.resolve([f3])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:f3
+      - fn*:f3()
+      - fn:bitcastF32toU32: Unsupported data types: vec2h. Supported types are: f32, vec2f, vec3f, vec4f.]
+    `);
+
+    const f4 = () => {
+      'use gpu';
+      const u = d.u32(1);
+      return std.bitcastF32toU32(u);
+    };
+    expect(() => tgpu.resolve([f4])).toThrowErrorMatchingInlineSnapshot(`
+      [Error: Resolution of the following tree failed:
+      - <root>
+      - fn*:f4
+      - fn*:f4()
+      - fn:bitcastF32toU32: Unsupported data types: u32. Supported types are: f32, vec2f, vec3f, vec4f.]
+    `);
   });
 });