Fix const fold

Author: l46kok

optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java

@@ -15,7 +15,6 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.collect.ImmutableList.toImmutableList;
-import static com.google.common.collect.MoreCollectors.onlyElement;
 import static dev.cel.checker.CelStandardDeclarations.StandardFunction.DURATION;
 import static dev.cel.checker.CelStandardDeclarations.StandardFunction.TIMESTAMP;
 
@@ -183,9 +182,9 @@ private boolean canFold(CelNavigableMutableExpr navigableExpr) {
         if (functionName.equals(Operator.EQUALS.getFunction())
             || functionName.equals(Operator.NOT_EQUALS.getFunction())) {
           if (mutableCall.args().stream()
-                  .anyMatch(node -> isExprConstantOfKind(node, CelConstant.Kind.BOOLEAN_VALUE))
+              .anyMatch(node -> isExprConstantOfKind(node, CelConstant.Kind.BOOLEAN_VALUE))
               || mutableCall.args().stream()
-                  .allMatch(node -> node.getKind().equals(Kind.CONSTANT))) {
+              .allMatch(node -> node.getKind().equals(Kind.CONSTANT))) {
             return true;
           }
         }
@@ -196,17 +195,69 @@ private boolean canFold(CelNavigableMutableExpr navigableExpr) {
 
         // Default case: all call arguments must be constants. If the argument is a container (ex:
         // list, map), then its arguments must be a constant.
-        return areChildrenArgConstant(navigableExpr);
+        return navigableExpr.children().allMatch(this::canEvaluate);
       case SELECT:
-        CelNavigableMutableExpr operand = navigableExpr.children().collect(onlyElement());
-        return areChildrenArgConstant(operand);
+        return navigableExpr.children().allMatch(this::canEvaluate);
       case COMPREHENSION:
-        return !isNestedComprehension(navigableExpr);
+        if (isNestedComprehension(navigableExpr)) {
+          return false;
+        }
+        CelMutableComprehension comprehension = navigableExpr.expr().comprehension();
+
+        if (!canEvaluate(CelNavigableMutableExpr.fromExpr(comprehension.iterRange()))
+            || !canEvaluate(CelNavigableMutableExpr.fromExpr(comprehension.accuInit()))) {
+          return false;
+        }
+
+        return canEvaluateComprehensionBody(CelNavigableMutableExpr.fromExpr(comprehension.loopStep()))
+            && canEvaluateComprehensionBody(CelNavigableMutableExpr.fromExpr(comprehension.loopCondition()));
       default:
         return false;
     }
   }
 
+  /**
+   * Checks if a subtree is safe to evaluate (i.e: it evaluates down to a constant expression)
+   */
+  private boolean canEvaluate(CelNavigableMutableExpr expression) {
+    return expression.allNodes().allMatch(this::isAllowedInConstantExpr);
+  }
+
+  private boolean canEvaluateComprehensionBody(CelNavigableMutableExpr expression) {
+    return expression.allNodes().allMatch(node -> {
+      Kind kind = node.getKind();
+      if (kind.equals(Kind.IDENT) || kind.equals(Kind.COMPREHENSION)) {
+        return true;
+      }
+      return isAllowedInConstantExpr(node);
+    });
+  }
+
+  private boolean isAllowedInConstantExpr(CelNavigableMutableExpr node) {
+    Kind kind = node.getKind();
+    if (kind.equals(Kind.CONSTANT)
+        || kind.equals(Kind.LIST)
+        || kind.equals(Kind.MAP)
+        || kind.equals(Kind.STRUCT)
+        || kind.equals(Kind.SELECT)) {
+      return true;
+    }
+    if (kind.equals(Kind.CALL)) {
+      CelMutableCall call = node.expr().call();
+      return foldableFunctions.contains(call.function());
+    }
+
+    return false;
+  }
+
+  private boolean isAllowedInFoldableExpr(CelNavigableMutableExpr node) {
+    Kind kind = node.getKind();
+    if (kind.equals(Kind.IDENT) || kind.equals(Kind.COMPREHENSION)) {
+      return true;
+    }
+    return isAllowedInConstantExpr(node);
+  }
+
   private boolean containsFoldableFunctionOnly(CelNavigableMutableExpr navigableExpr) {
     return navigableExpr
         .allNodes()
@@ -248,22 +299,6 @@ private static boolean canFoldInOperator(CelNavigableMutableExpr navigableExpr)
     return true;
   }
 
-  private static boolean areChildrenArgConstant(CelNavigableMutableExpr expr) {
-    if (expr.getKind().equals(Kind.CONSTANT)) {
-      return true;
-    }
-
-    if (expr.getKind().equals(Kind.CALL)
-        || expr.getKind().equals(Kind.LIST)
-        || expr.getKind().equals(Kind.MAP)
-        || expr.getKind().equals(Kind.SELECT)
-        || expr.getKind().equals(Kind.STRUCT)) {
-      return expr.children().allMatch(ConstantFoldingOptimizer::areChildrenArgConstant);
-    }
-
-    return false;
-  }
-
   private static boolean isNestedComprehension(CelNavigableMutableExpr expr) {
     Optional<CelNavigableMutableExpr> maybeParent = expr.parent();
     while (maybeParent.isPresent()) {

optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java

@@ -48,6 +48,7 @@
 public class ConstantFoldingOptimizerTest {
   private static final CelOptions CEL_OPTIONS =
       CelOptions.current()
+          .populateMacroCalls(true)
           .enableTimestampEpoch(true)
           .build();
   private static final Cel CEL =
@@ -56,12 +57,23 @@ public class ConstantFoldingOptimizerTest {
           .addVar("y", SimpleType.DYN)
           .addVar("list_var", ListType.create(SimpleType.STRING))
           .addVar("map_var", MapType.create(SimpleType.STRING, SimpleType.STRING))
+          .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
           .addFunctionDeclarations(
               CelFunctionDecl.newFunctionDeclaration(
                   "get_true",
-                  CelOverloadDecl.newGlobalOverload("get_true_overload", SimpleType.BOOL)))
+                  CelOverloadDecl.newGlobalOverload("get_true_overload", SimpleType.BOOL)),
+              CelFunctionDecl.newFunctionDeclaration(
+                  "get_list",
+                  CelOverloadDecl.newGlobalOverload(
+                      "get_list_overload",
+                      ListType.create(SimpleType.INT),
+                      ListType.create(SimpleType.INT)))
+          )
           .addFunctionBindings(
-              CelFunctionBinding.from("get_true_overload", ImmutableList.of(), unused -> true))
+              CelFunctionBinding.from("get_true_overload", ImmutableList.of(), unused -> true),
+              CelFunctionBinding.from(
+                  "get_list_overload", ImmutableList.class, arg -> arg)
+          )
           .addMessageTypes(TestAllTypes.getDescriptor())
           .setContainer(CelContainer.ofName("cel.expr.conformance.proto3"))
           .setOptions(CEL_OPTIONS)
@@ -371,6 +383,8 @@ public void constantFold_macros_withoutMacroCallMetadata(String source) throws E
   @TestParameters("{source: 'x == 42'}")
   @TestParameters("{source: 'timestamp(100)'}")
   @TestParameters("{source: 'duration(\"1h\")'}")
+  @TestParameters("{source: '[true].exists(x, x == get_true())'}")
+  @TestParameters("{source: 'get_list([1, 2]).map(x, x * 2)'}")
   public void constantFold_noOp(String source) throws Exception {
     CelAbstractSyntaxTree ast = CEL.compile(source).getAst();
 

tools/src/test/java/dev/cel/tools/ai/AgenticPolicyCompilerTest.java

@@ -283,6 +283,7 @@ private void runTests(Cel cel, CelAbstractSyntaxTree ast, PolicyTestSuite testSu
                   ? (List<AgentMessage>) inputMap.get("_test_history")
                   : ImmutableList.of();
 
+          @SuppressWarnings("Immutable")
           CelLateFunctionBindings bindings = CelLateFunctionBindings.from(
               CelFunctionBinding.from(
                   "agent_history",