feat: bitcastF32toU32

packages/typegpu/src/data/numberOps.ts

@@ -27,3 +27,9 @@ export function bitcastU32toI32Impl(n: number): number {
   dataView.setUint32(0, n, true);
   return dataView.getInt32(0, true);
 }
+
+export function bitcastF32toU32Impl(n: number): number {
+  const dataView = new DataView(new ArrayBuffer(4));
+  dataView.setFloat32(0, n, true);
+  return dataView.getUint32(0, true);
+}

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,9 +1,13 @@
 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';
@@ -65,3 +69,36 @@ export const bitcastU32toI32 = dualImpl({
     };
   },
 });
+
+export type BitcastF32toU32Overload = ((value: number) => number) &
+  ((value: v2f) => v2u) &
+  ((value: v3f) => v3u) &
+  ((value: v4f) => v4u);
+
+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 = unify(arg, [f32]) ?? arg;
+    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,6 +186,6 @@ 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/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);
@@ -128,6 +158,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 {
@@ -139,11 +170,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 {
@@ -155,5 +192,10 @@ 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);
+      }"
+    `);
   });
 });