From 5905379519ff8b44f10a431c25bf4a9cf6f3933a Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Fri, 19 Sep 2025 12:09:00 -0700
Subject: [PATCH] Allow constant folding to fold equals operator

---
 .../optimizers/ConstantFoldingOptimizer.java  | 29 +++++++++++++++++++
 .../ConstantFoldingOptimizerTest.java         |  8 +++++
 2 files changed, 37 insertions(+)

diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
index fd52f4138..a2c94de40 100644
--- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
+++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java
@@ -167,6 +167,10 @@ private boolean canFold(CelNavigableMutableExpr navigableExpr) {
               && cond.constant().getKind().equals(CelConstant.Kind.BOOLEAN_VALUE);
         }
 
+        if (functionName.equals(Operator.EQUALS.getFunction())) {
+          return mutableCall.args().stream().anyMatch(x -> isExprConstantOfKind(x, CelConstant.Kind.BOOLEAN_VALUE));
+        }
+
         if (functionName.equals(Operator.IN.getFunction())) {
           return canFoldInOperator(navigableExpr);
         }
@@ -393,6 +397,27 @@ private Optional<CelMutableAst> maybePruneBranches(
           }
         }
       }
+    } else if (function.equals(Operator.EQUALS.getFunction())) {
+      CelMutableExpr lhs = call.args().get(0);
+      CelMutableExpr rhs = call.args().get(1);
+      boolean lhsIsBoolean = isExprConstantOfKind(lhs, CelConstant.Kind.BOOLEAN_VALUE);
+      boolean rhsIsBoolean = isExprConstantOfKind(rhs, CelConstant.Kind.BOOLEAN_VALUE);
+      Optional<CelMutableExpr> branchToRetain = Optional.empty();
+
+      if (lhsIsBoolean && rhsIsBoolean) {
+        boolean result = lhs.constant().booleanValue() == rhs.constant().booleanValue();
+        branchToRetain = Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(result)));
+      } else if (lhsIsBoolean) {
+        branchToRetain = Optional.of(lhs.constant().booleanValue() ? rhs : lhs);
+      } else if (rhsIsBoolean) {
+        branchToRetain = Optional.of(rhs.constant().booleanValue() ? lhs : rhs);
+      }
+
+      return branchToRetain.map(node ->
+          astMutator.replaceSubtree(
+              mutableAst,
+              node,
+              expr.id()));
     }
 
     return Optional.empty();
@@ -663,6 +688,10 @@ public static Builder newBuilder() {
     ConstantFoldingOptions() {}
   }
 
+  private static boolean isExprConstantOfKind(CelMutableExpr expr, CelConstant.Kind constantKind) {
+    return expr.getKind().equals(Kind.CONSTANT) && expr.constant().getKind().equals(constantKind);
+  }
+
   private ConstantFoldingOptimizer(ConstantFoldingOptions constantFoldingOptions) {
     this.constantFoldingOptions = constantFoldingOptions;
     this.astMutator = AstMutator.newInstance(constantFoldingOptions.maxIterationLimit());
diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
index ea1e77edb..acd0770f7 100644
--- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
+++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java
@@ -189,6 +189,14 @@ public class ConstantFoldingOptimizerTest {
   @TestParameters("{source: 'sets.contains([1], [1])', expected: 'true'}")
   @TestParameters(
       "{source: 'cel.bind(r0, [1, 2, 3], cel.bind(r1, 1 in r0, r1))', expected: 'true'}")
+  @TestParameters("{source: 'x == true', expected: 'x'}")
+  @TestParameters("{source: 'true == x', expected: 'x'}")
+  @TestParameters("{source: 'x == false', expected: 'false'}")
+  @TestParameters("{source: 'false == x', expected: 'false'}")
+  @TestParameters("{source: 'true == false', expected: 'false'}")
+  @TestParameters("{source: 'true == true', expected: 'true'}")
+  @TestParameters("{source: 'false == true', expected: 'false'}")
+  @TestParameters("{source: 'false == false', expected: 'true'}")
   // TODO: Support folding lists with mixed types. This requires mutable lists.
   // @TestParameters("{source: 'dyn([1]) + [1.0]'}")
   public void constantFold_success(String source, String expected) throws Exception {