initial short circuit eval

Author: cieplypolar

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -352,6 +352,33 @@ ${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)) {
+        let evalRhs = !lhsExpr.value;
+        if (op === '&&') {
+          evalRhs = !evalRhs;
+        }
+
+        if (!evalRhs) {
+          return op === '||' ? snip(true, bool, 'constant') : snip(false, bool, 'constant');
+        }
+
+        const rhsExpr = this._expression(rhs);
+
+        if (isKnownAtComptime(rhsExpr)) {
+          return snip(rhsExpr.value, bool, 'constant');
+        }
+
+        if (rhsExpr.dataType === UnknownData) {
+          throw new WgslTypeError(`Right-hand side of '${op}' is of unknown type`);
+        }
+
+        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

@@ -123,9 +123,6 @@ describe('not', () => {
 
     expect(tgpu.resolve([f])).toMatchInlineSnapshot(`
       "fn f() -> i32 {
-        if (((false && true) && false)) {
-          return 1;
-        }
         return -1;
       }"
     `);

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

@@ -2137,13 +2137,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', () => {
@@ -2164,13 +2164,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 `!`', () => {
@@ -2183,10 +2183,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', () => {
@@ -2207,4 +2207,127 @@ 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 `||` short-circuit evaluation', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        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 `&&` short-circuit evaluation', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        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 `||` short-circuit evaluation', () => {
+      state.result = false;
+
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        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 `&&` short-circuit evaluation', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        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 short-circuit evaluation', () => {
+      const f = () => {
+        'use gpu';
+        let res = -1;
+        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);
+    });
+  });
 });