correct NaN handling

Author: cieplypolar

packages/typegpu/src/std/boolean.ts

@@ -186,12 +186,14 @@ function cpuNot(value: unknown): boolean | AnyBooleanVecInstance {
 
   if (isVecInstance(value)) {
     if (value.length === 2) {
-      return vec2b(!value.x, !value.y);
+      return vec2b(cpuNot(value.x), cpuNot(value.y));
     }
     if (value.length === 3) {
-      return vec3b(!value.x, !value.y, !value.z);
+      return vec3b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z));
+    }
+    if (value.length === 4) {
+      return vec4b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z), cpuNot(value.w));
     }
-    return vec4b(!value.x, !value.y, !value.z, !value.w);
   }
 
   return !value;
@@ -224,10 +226,7 @@ export const not = dualImpl({
   normalImpl: cpuNot,
   codegenImpl: (_ctx, [arg]) => {
     if (isKnownAtComptime(arg)) {
-      if (typeof arg.value === 'number' && isNaN(arg.value)) {
-        return 'false';
-      }
-      return `${!arg.value}`;
+      return `${cpuNot(arg.value)}`;
     }
 
     const { dataType } = arg;

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -173,8 +173,15 @@ const unaryOpCodeToCodegen = {
       return snip(`!${argStr}`, bool, 'runtime');
     }
     if (wgsl.isNumericSchema(dataType)) {
-      // no bool cast, because bool(NaN) is true in WGSL but in JS it's false
-      return snip(`!(${argStr} != 0)`, bool, 'runtime');
+      const resultStr = `!bool(${argStr})`;
+      const nanGuardedStr = // abstractFloat will be resolved as comptime
+        dataType.type === 'f32'
+          ? `(((bitcast<u32>(${argStr}) & 0x7fffffff) > 0x7f800000) || ${resultStr})`
+          : dataType.type === 'f16'
+            ? `(((bitcast<u32>(${argStr}) & 0x7fff) > 0x7c00) || ${resultStr})`
+            : resultStr;
+
+      return snip(nanGuardedStr, bool, 'runtime');
     }
 
     return snip(false, bool, 'constant');

packages/typegpu/tests/std/boolean/not.test.ts

@@ -96,6 +96,19 @@ describe('not', () => {
     `);
   });
 
+  it('generates correct WGSL on a numeric vector comptime-known argument', () => {
+    const f = () => {
+      'use gpu';
+      const v = not(d.vec4f(Infinity, -Infinity, 0, NaN));
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() {
+        var v = vec4<bool>(false, false, true, false);
+      }"
+    `);
+  });
+
   it('evaluates at compile time for comptime-known arguments', () => {
     const getN = tgpu.comptime(() => 42);
     const slot = tgpu.slot<{ a?: number }>({});