feat: unary operator ! support and std.not (#2346)

Author: cieplypolar

apps/typegpu-docs/tests/individual-example-tests/tgsl-parsing-test.test.ts

@@ -44,11 +44,8 @@ describe('tgsl parsing test example', () => {
         s = (s && true);
         s = (s && true);
         s = (s && true);
-        s = (s && !false);
         s = (s && true);
-        s = (s && !false);
         s = (s && true);
-        s = (s && !false);
         s = (s && true);
         s = (s && true);
         s = (s && true);
@@ -58,9 +55,12 @@ describe('tgsl parsing test example', () => {
         s = (s && true);
         s = (s && true);
         s = (s && true);
-        s = (s && !false);
         s = (s && true);
-        s = (s && !false);
+        s = (s && true);
+        s = (s && true);
+        s = (s && true);
+        s = (s && true);
+        s = (s && true);
         s = (s && true);
         s = (s && true);
         var vec = vec3<bool>(true, false, true);

packages/typegpu/src/data/wgslTypes.ts

@@ -1713,6 +1713,10 @@ export function isVoid(value: unknown): value is Void {
   return isMarkedInternal(value) && (value as Void).type === 'void';
 }
 
+export function isBool(value: unknown): value is Bool {
+  return isMarkedInternal(value) && (value as Bool).type === 'bool';
+}
+
 export function isNumericSchema(
   schema: unknown,
 ): schema is AbstractInt | AbstractFloat | F32 | F16 | I32 | U32 {

packages/typegpu/src/std/boolean.ts

@@ -13,6 +13,10 @@ import {
   type AnyVecInstance,
   type AnyWgslData,
   type BaseData,
+  isBool,
+  isNumericSchema,
+  isVec,
+  isVecBool,
   isVecInstance,
   type v2b,
   type v3b,
@@ -164,19 +168,82 @@ export const ge = dualImpl({
 
 // logical ops
 
-const cpuNot = <T extends AnyBooleanVecInstance>(value: T): T => VectorOps.neg[value.kind](value);
+type VecInstanceToBooleanVecInstance<T extends AnyVecInstance> = T extends AnyVec2Instance
+  ? v2b
+  : T extends AnyVec3Instance
+    ? v3b
+    : v4b;
+
+function cpuNot(value: boolean): boolean;
+function cpuNot(value: number): boolean;
+function cpuNot<T extends AnyVecInstance>(value: T): VecInstanceToBooleanVecInstance<T>;
+function cpuNot(value: unknown): boolean;
+function cpuNot(value: unknown): boolean | AnyBooleanVecInstance {
+  if (typeof value === 'number' && isNaN(value)) {
+    return false;
+  }
+
+  if (isVecInstance(value)) {
+    if (value.length === 2) {
+      return vec2b(cpuNot(value.x), cpuNot(value.y));
+    }
+    if (value.length === 3) {
+      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 !value;
+}
 
 /**
- * Returns **component-wise** `!value`.
+ * Returns the logical negation of the given value.
+ * For scalars (bool, number), returns `!value`.
+ * For boolean vectors, returns **component-wise** `!value`.
+ * For numeric vectors, returns a boolean vector with component-wise truthiness negation.
+ * For all other types, returns the truthiness negation (in WGSL, this applies only if the value is known at compile-time).
  * @example
- * not(vec2b(false, true)) // returns vec2b(true, false)
+ * not(true) // returns false
+ * not(-1) // returns false
+ * not(0) // returns true
  * not(vec3b(true, true, false)) // returns vec3b(false, false, true)
+ * not(vec3f(1.0, 0.0, -1.0)) // returns vec3b(false, true, false)
+ * not({a: 1882}) // returns false
+ * not(NaN) // returns false **as in WGSL**
  */
 export const not = dualImpl({
   name: 'not',
-  signature: (...argTypes) => ({ argTypes, returnType: argTypes[0] }),
+  signature: (arg) => {
+    const returnType = isVec(arg) ? correspondingBooleanVectorSchema(arg) : bool;
+    return {
+      argTypes: [arg],
+      returnType,
+    };
+  },
   normalImpl: cpuNot,
-  codegenImpl: (_ctx, [arg]) => stitch`!(${arg})`,
+  codegenImpl: (_ctx, [arg]) => {
+    const { dataType } = arg;
+
+    if (isBool(dataType)) {
+      return stitch`!${arg}`;
+    }
+    if (isNumericSchema(dataType)) {
+      return stitch`!bool(${arg})`;
+    }
+
+    if (isVecBool(dataType)) {
+      return stitch`!(${arg})`;
+    }
+
+    if (isVec(dataType)) {
+      const vecConstructorStr = `vec${dataType.componentCount}<bool>`;
+      return stitch`!(${vecConstructorStr}(${arg}))`;
+    }
+
+    return 'false';
+  },
 });
 
 const cpuOr = <T extends AnyBooleanVecInstance>(lhs: T, rhs: T) => VectorOps.or[lhs.kind](lhs, rhs);

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -157,6 +157,35 @@ function operatorToType<
 const unaryOpCodeToCodegen = {
   '-': neg[$gpuCallable].call.bind(neg),
   void: () => snip(undefined, wgsl.Void, 'constant'),
+  '!': (ctx: GenerationCtx, [argExpr]: Snippet[]) => {
+    if (argExpr === undefined) {
+      throw new Error('The unary operator `!` expects 1 argument, but 0 were provided.');
+    }
+
+    if (isKnownAtComptime(argExpr)) {
+      return snip(!argExpr.value, bool, 'constant');
+    }
+
+    const { value, dataType } = argExpr;
+    const argStr = ctx.resolve(value, dataType).value;
+
+    if (wgsl.isBool(dataType)) {
+      return snip(`!${argStr}`, bool, 'runtime');
+    }
+    if (wgsl.isNumericSchema(dataType)) {
+      const resultStr = `!bool(${argStr})`;
+      const nanGuardedStr = // abstractFloat will be resolved as comptime known value
+        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');
+  },
 } satisfies Partial<Record<tinyest.UnaryOperator, (...args: never[]) => unknown>>;
 
 const binaryOpCodeToCodegen = {

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

@@ -1,11 +1,156 @@
-import { describe, expect, it } from 'vitest';
-import { vec2b, vec3b, vec4b } from '../../../src/data/index.ts';
+import { describe, expect } from 'vitest';
+import { it } from 'typegpu-testing-utility';
 import { not } from '../../../src/std/boolean.ts';
+import tgpu, { d, std } from '../../../src/index.js';
 
-describe('neg', () => {
-  it('negates', () => {
-    expect(not(vec2b(true, false))).toStrictEqual(vec2b(false, true));
-    expect(not(vec3b(false, false, true))).toStrictEqual(vec3b(true, true, false));
-    expect(not(vec4b(true, true, false, false))).toStrictEqual(vec4b(false, false, true, true));
+describe('not', () => {
+  it('negates booleans', () => {
+    expect(not(true)).toBe(false);
+    expect(not(false)).toBe(true);
+  });
+
+  it('converts numbers to booleans and negates', () => {
+    expect(not(0)).toBe(true);
+    expect(not(-1)).toBe(false);
+    expect(not(42)).toBe(false);
+  });
+
+  it('negates boolean vectors', () => {
+    expect(not(d.vec2b(true, false))).toStrictEqual(d.vec2b(false, true));
+    expect(not(d.vec3b(false, false, true))).toStrictEqual(d.vec3b(true, true, false));
+    expect(not(d.vec4b(true, true, false, false))).toStrictEqual(d.vec4b(false, false, true, true));
+  });
+
+  it('converts numeric vectors to booleans vectors and negates component-wise', () => {
+    expect(not(d.vec2f(0.0, 1.0))).toStrictEqual(d.vec2b(true, false));
+    expect(not(d.vec3i(0, 5, -1))).toStrictEqual(d.vec3b(true, false, false));
+    expect(not(d.vec4u(0, 0, 1, 0))).toStrictEqual(d.vec4b(true, true, false, true));
+    expect(not(d.vec4h(0, 3.14, 0, -2.5))).toStrictEqual(d.vec4b(true, false, true, false));
+  });
+
+  it('negates truthiness check', () => {
+    expect(not(null)).toBe(true);
+    expect(not(undefined)).toBe(true);
+    expect(not({})).toBe(false);
+  });
+
+  it('mimics WGSL behavior on NaN', () => {
+    expect(not(NaN)).toBe(false);
+  });
+
+  it('generates correct WGSL on a boolean runtime-known argument', () => {
+    const testFn = tgpu.fn(
+      [d.bool],
+      d.bool,
+    )((v) => {
+      return not(v);
+    });
+    expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(`
+      "fn testFn(v: bool) -> bool {
+        return !v;
+      }"
+    `);
+  });
+
+  it('generates correct WGSL on a numeric runtime-known argument', () => {
+    const testFn = tgpu.fn(
+      [d.i32],
+      d.bool,
+    )((v) => {
+      return not(v);
+    });
+    expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(`
+        "fn testFn(v: i32) -> bool {
+          return !bool(v);
+        }"
+      `);
+  });
+
+  it('generates correct WGSL on a boolean vector runtime-known argument', () => {
+    const testFn = tgpu.fn(
+      [d.vec3b],
+      d.vec3b,
+    )((v) => {
+      return not(v);
+    });
+    expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(`
+      "fn testFn(v: vec3<bool>) -> vec3<bool> {
+        return !(v);
+      }"
+    `);
+  });
+
+  it('generates correct WGSL on a numeric vector runtime-known argument', () => {
+    const testFn = tgpu.fn(
+      [d.vec3f],
+      d.vec3b,
+    )((v) => {
+      return not(v);
+    });
+    expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(`
+      "fn testFn(v: vec3f) -> vec3<bool> {
+        return !(vec3<bool>(v));
+      }"
+    `);
+  });
+
+  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 }>({});
+
+    const f = () => {
+      'use gpu';
+      if (not(getN()) && not(slot.$.a) && not(d.vec4f(1, 8, 8, 2)).x) {
+        return 1;
+      }
+      return -1;
+    };
+
+    expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+      "fn f() -> i32 {
+        if (((false && true) && false)) {
+          return 1;
+        }
+        return -1;
+      }"
+    `);
+  });
+
+  it('mimics JS on non-primitive values', ({ root }) => {
+    const buffer = root.createUniform(d.mat4x4f);
+    const testFn = tgpu.fn([d.vec3f, d.atomic(d.u32), d.ptrPrivate(d.u32)])((v, a, p) => {
+      const _b0 = not(buffer);
+      const _b1 = not(buffer.$);
+      const _b2 = not(v);
+      const _b3 = not(a);
+      const _b4 = not(std.atomicLoad(a));
+      const _b5 = not(p);
+      const _b6 = not(p.$);
+    });
+
+    expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(`
+      "fn testFn(v: vec3f, a: atomic<u32>, p: ptr<private, u32>) {
+        const _b0 = false;
+        let _b1 = false;
+        var _b2 = !(vec3<bool>(v));
+        let _b3 = false;
+        let _b4 = !bool(atomicLoad(&a));
+        let _b5 = false;
+        let _b6 = !bool((*p));
+      }"
+    `);
   });
 });