Internal Change

PiperOrigin-RevId: 934764659

Author: CEL Dev Team

optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel

@@ -28,6 +28,7 @@ java_library(
         "//common/ast",
         "//common/ast:mutable_expr",
         "//common/internal:date_time_helpers",
+        "//common/navigation:common",
         "//common/navigation:mutable_navigation",
         "//common/types",
         "//extensions:optional_library",

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

@@ -41,6 +41,7 @@
 import dev.cel.common.internal.DateTimeHelpers;
 import dev.cel.common.navigation.CelNavigableMutableAst;
 import dev.cel.common.navigation.CelNavigableMutableExpr;
+import dev.cel.common.navigation.TraversalOrder;
 import dev.cel.common.types.SimpleType;
 import dev.cel.extensions.CelOptionalLibrary.Function;
 import dev.cel.optimizer.AstMutator;
@@ -111,7 +112,7 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
       ImmutableList<CelNavigableMutableExpr> foldableExprs =
           CelNavigableMutableAst.fromAst(mutableAst)
               .getRoot()
-              .allNodes()
+              .allNodes(TraversalOrder.PRE_ORDER)
               .filter(this::canFold)
               .collect(toImmutableList());
       for (CelNavigableMutableExpr foldableExpr : foldableExprs) {
@@ -122,7 +123,13 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
         mutatedResult = maybePruneBranches(mutableAst, foldableExpr.expr());
         if (!mutatedResult.isPresent()) {
           // Evaluate the call then fold
-          mutatedResult = maybeFold(optimizerEnv, mutableAst, foldableExpr);
+          try {
+            mutatedResult = maybeFold(optimizerEnv, mutableAst, foldableExpr);
+          } catch (CelEvaluationException e) {
+            throw new CelOptimizationException(
+                "Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(),
+                e);
+          }
         }
 
         if (!mutatedResult.isPresent()) {
@@ -132,12 +139,17 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel)
 
         continueFolding = true;
         mutableAst = mutatedResult.get();
+        // Break the loop because we mutated the AST. Since we traverse in PRE_ORDER (top-down),
+        // mutating a parent node means its children are now obsolete or folded.
+        // We restart the traversal to gather a fresh list of foldable expressions.
+        break;
       }
     }
 
     // If the output is a list, map, or struct which contains optional entries, then prune it
     // to make sure that the optionals, if resolved, do not surface in the output literal.
     mutableAst = pruneOptionalElements(mutableAst);
+
     return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst());
   }
 
@@ -280,11 +292,11 @@ private static boolean isNestedComprehension(CelNavigableMutableExpr expr) {
 
   private Optional<CelMutableAst> maybeFold(
       Cel cel, CelMutableAst mutableAst, CelNavigableMutableExpr node)
-      throws CelOptimizationException {
+      throws CelOptimizationException, CelEvaluationException {
     Object result;
     try {
       result = evaluateExpr(cel, node);
-    } catch (CelValidationException | CelEvaluationException e) {
+    } catch (CelValidationException e) {
       throw new CelOptimizationException(
           "Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(), e);
     }

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

@@ -252,6 +252,9 @@ private static Cel setupEnv(CelBuilder celBuilder) {
   @TestParameters(
       "{source: 'timestamp(\"2000-01-01T00:02:03.2123Z\") + duration(\"25h2m32s42ms53us29ns\")',"
           + " expected: 'timestamp(\"2000-01-02T01:04:35.254353029Z\")'}")
+  @TestParameters(
+      "{source: 'has({\"req\": \"Avail\"}.opt) ? ({\"req\": \"Avail\"}.req + \" \" +"
+          + " {\"req\": \"Avail\"}.opt) : {\"req\": \"Avail\"}.req', expected: '\"Avail\"'}")
   // 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 {
@@ -534,26 +537,20 @@ public void constantFold_astProducesConsistentlyNumberedIds() throws Exception {
 
   @Test
   public void iterationLimitReached_throws() throws Exception {
-    StringBuilder sb = new StringBuilder();
-    sb.append("0");
-    for (int i = 1; i < 200; i++) {
-      sb.append(" + ").append(i);
-    } // 0 + 1 + 2 + 3 + ... 200
     Cel cel =
         runtimeFlavor
             .builder()
             .setOptions(
                 CelOptions.current()
                     .enableHeterogeneousNumericComparisons(true)
-                    .maxParseRecursionDepth(200)
                     .build())
             .build();
-    CelAbstractSyntaxTree ast = cel.compile(sb.toString()).getAst();
+    CelAbstractSyntaxTree ast = cel.compile("1 + 1").getAst();
     CelOptimizer optimizer =
         CelOptimizerFactory.standardCelOptimizerBuilder(cel)
             .addAstOptimizers(
                 ConstantFoldingOptimizer.newInstance(
-                    ConstantFoldingOptions.newBuilder().maxIterationLimit(200).build()))
+                    ConstantFoldingOptions.newBuilder().maxIterationLimit(1).build()))
             .build();
 
     CelOptimizationException e =