impr: Short-circuit evaluation for operators && and ||

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -352,6 +352,31 @@ ${this.ctx.pre}}`;
       // Logical/Binary/Assignment Expression
       const [exprType, lhs, op, rhs] = expression;
       const lhsExpr = this._expression(lhs);
+
+      // Short Circuit Evaluation
+      if ((op === '||' || op === '&&') && isKnownAtComptime(lhsExpr)) {
+        const evalRhs = op === '&&' ? lhsExpr.value : !lhsExpr.value;
+
+        if (!evalRhs) {
+          return snip(op === '||', bool, 'constant');
+        }
+
+        const rhsExpr = this._expression(rhs);
+
+        if (rhsExpr.dataType === UnknownData) {
+          throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
+        }
+
+        if (isKnownAtComptime(rhsExpr)) {
+          return snip(!!rhsExpr.value, bool, 'constant');
+        }
+
+        // we can skip lhs
+        const convRhs = tryConvertSnippet(this.ctx, rhsExpr, bool, false);
+        const rhsStr = this.ctx.resolve(convRhs.value, convRhs.dataType).value;
+        return snip(rhsStr, bool, 'runtime');
+      }
+
       const rhsExpr = this._expression(rhs);
 
       if (rhsExpr.value instanceof RefOperator) {

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

@@ -121,9 +121,6 @@ describe('not', () => {
 
     expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
       "fn f() -> i32 {
-        if (((false && true) && false)) {
-          return 1;
-        }
         return -1;
       }"
     `);
@@ -153,4 +150,110 @@ describe('not', () => {
       }"
     `);
   });
+
+  it('converts numeric vectors to boolean 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', () => {
+    const s = {};
+    expect(not(null)).toBe(true);
+    expect(not(undefined)).toBe(true);
+    expect(not(s)).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 {
+        return -1;
+      }"
+    `);
+  });
 });

packages/typegpu/tests/tgsl/wgslGenerator.test.ts

@@ -2139,13 +2139,13 @@ describe('wgslGenerator', () => {
     };
 
     expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
-        "fn f() -> i32 {
-          if ((true || false)) {
-            return 1;
-          }
-          return -1;
-        }"
-      `);
+      "fn f() -> i32 {
+        {
+          return 1;
+        }
+        return -1;
+      }"
+    `);
   });
 
   it('handles unary operator `!` on operands from slots and accessors', () => {
@@ -2166,13 +2166,13 @@ describe('wgslGenerator', () => {
     };
 
     expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
-        "fn f() -> i32 {
-          if ((true && true)) {
-            return 1;
-          }
-          return -1;
-        }"
-      `);
+      "fn f() -> i32 {
+        {
+          return 1;
+        }
+        return -1;
+      }"
+    `);
   });
 
   it('handles chained unary operators `!`', () => {
@@ -2185,10 +2185,10 @@ describe('wgslGenerator', () => {
     });
 
     expect(tgpu.resolve([testFn])).toMatchInlineSnapshot(`
-        "fn testFn(n: i32) -> bool {
-          return (true || !!!bool(n));
-        }"
-      `);
+      "fn testFn(n: i32) -> bool {
+        return true;
+      }"
+    `);
   });
 
   it('handles unary operator `!` on complex comptime-known operand', () => {
@@ -2209,4 +2209,178 @@ describe('wgslGenerator', () => {
         }"
       `);
   });
+
+  describe('short-circuit evaluation', () => {
+    const state = {
+      counter: 0,
+      result: true,
+    };
+
+    const getTrackedBool = tgpu.comptime(() => {
+      state.counter++;
+      return state.result;
+    });
+
+    beforeEach(() => {
+      state.counter = 0;
+      state.result = true;
+    });
+
+    it('handles `||`', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test
+        if (true || getTrackedBool()) {
+          res = 1;
+        }
+        return res;
+      };
+
+      expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+        "fn f() -> i32 {
+          var res = -1;
+          {
+            res = 1i;
+          }
+          return res;
+        }"
+      `);
+      expect(state.counter).toBe(0);
+    });
+
+    it('handles `&&`', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test
+        if (false && getTrackedBool()) {
+          res = 1;
+        }
+        return res;
+      };
+
+      expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+        "fn f() -> i32 {
+          var res = -1;
+          return res;
+        }"
+      `);
+      expect(state.counter).toBe(0);
+    });
+
+    it('handles chained `||`', () => {
+      state.result = false;
+
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test
+        if (getTrackedBool() || true || getTrackedBool() || getTrackedBool() || getTrackedBool()) {
+          res = 1;
+        }
+        return res;
+      };
+
+      expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+        "fn f() -> i32 {
+          var res = -1;
+          {
+            res = 1i;
+          }
+          return res;
+        }"
+      `);
+      expect(state.counter).toEqual(1);
+    });
+
+    it('handles chained `&&`', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test
+        if (getTrackedBool() && false && getTrackedBool() && getTrackedBool() && getTrackedBool()) {
+          res = 1;
+        }
+        return res;
+      };
+
+      expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+        "fn f() -> i32 {
+          var res = -1;
+          return res;
+        }"
+      `);
+      expect(state.counter).toBe(1);
+    });
+
+    it('handles mixed logical operators', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test
+        if (true || (getTrackedBool() && getTrackedBool())) {
+          res = 1;
+        }
+        return res;
+      };
+
+      expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
+        "fn f() -> i32 {
+          var res = -1;
+          {
+            res = 1i;
+          }
+          return res;
+        }"
+      `);
+      expect(state.counter).toBe(0);
+    });
+
+    it('skips lhs if known at compile time', () => {
+      const f1 = tgpu.fn(
+        [d.bool],
+        d.i32,
+      )((b) => {
+        'use gpu';
+        let res = -1;
+        // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test
+        if (false || b) {
+          res = 1;
+        }
+        return res;
+      });
+
+      const f2 = tgpu.fn(
+        [d.bool],
+        d.i32,
+      )((b) => {
+        'use gpu';
+        let res = -1;
+        // oxlint-disable-next-line(no-constant-binary-expression) -- part of the test
+        if (true && b) {
+          res = 1;
+        }
+        return res;
+      });
+
+      expect(tgpu.resolve([f1, f2])).toMatchInlineSnapshot(`
+        "fn f1(b: bool) -> i32 {
+          var res = -1;
+          if (b) {
+            res = 1i;
+          }
+          return res;
+        }
+
+        fn f2(b: bool) -> i32 {
+          var res = -1;
+          if (b) {
+            res = 1i;
+          }
+          return res;
+        }"
+      `);
+    });
+  });
 });