Optimize composed policies using Constant Folding and Common Subexpression Elimination

PiperOrigin-RevId: 797449393

Author: l46kok

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

@@ -202,7 +202,9 @@ private static boolean canFoldInOperator(CelNavigableMutableExpr navigableExpr)
       CelNavigableMutableExpr parent = identNode.parent().orElse(null);
       while (parent != null) {
         if (parent.getKind().equals(Kind.COMPREHENSION)) {
-          if (parent.expr().comprehension().accuVar().equals(identNode.expr().ident().name())) {
+          String identName = identNode.expr().ident().name();
+          if (parent.expr().comprehension().accuVar().equals(identName)
+              || parent.expr().comprehension().iterVar().equals(identName)) {
             // Prevent folding a subexpression if it contains a variable declared by a
             // comprehension. The subexpression cannot be compiled without the full context of the
             // surrounding comprehension.

parser/src/main/java/dev/cel/parser/Operator.java

@@ -103,22 +103,34 @@ static Optional<Operator> find(String text) {
   private static final ImmutableMap<String, Operator> REVERSE_OPERATORS =
       ImmutableMap.<String, Operator>builder()
           .put(ADD.getFunction(), ADD)
+          .put(ALL.getFunction(), ALL)
+          .put(CONDITIONAL.getFunction(), CONDITIONAL)
           .put(DIVIDE.getFunction(), DIVIDE)
           .put(EQUALS.getFunction(), EQUALS)
+          .put(EXISTS.getFunction(), EXISTS)
+          .put(EXISTS_ONE.getFunction(), EXISTS_ONE)
+          .put(FILTER.getFunction(), FILTER)
           .put(GREATER.getFunction(), GREATER)
           .put(GREATER_EQUALS.getFunction(), GREATER_EQUALS)
+          .put(HAS.getFunction(), HAS)
           .put(IN.getFunction(), IN)
+          .put(INDEX.getFunction(), INDEX)
           .put(LESS.getFunction(), LESS)
           .put(LESS_EQUALS.getFunction(), LESS_EQUALS)
           .put(LOGICAL_AND.getFunction(), LOGICAL_AND)
           .put(LOGICAL_NOT.getFunction(), LOGICAL_NOT)
           .put(LOGICAL_OR.getFunction(), LOGICAL_OR)
+          .put(MAP.getFunction(), MAP)
           .put(MODULO.getFunction(), MODULO)
           .put(MULTIPLY.getFunction(), MULTIPLY)
           .put(NEGATE.getFunction(), NEGATE)
           .put(NOT_EQUALS.getFunction(), NOT_EQUALS)
-          .put(SUBTRACT.getFunction(), SUBTRACT)
+          .put(NOT_STRICTLY_FALSE.getFunction(), NOT_STRICTLY_FALSE)
           .put(OLD_IN.getFunction(), OLD_IN)
+          .put(OLD_NOT_STRICTLY_FALSE.getFunction(), OLD_NOT_STRICTLY_FALSE)
+          .put(OPTIONAL_INDEX.getFunction(), OPTIONAL_INDEX)
+          .put(OPTIONAL_SELECT.getFunction(), OPTIONAL_SELECT)
+          .put(SUBTRACT.getFunction(), SUBTRACT)
           .buildOrThrow();
 
   // precedence of the operator, where the higher value means higher.
@@ -168,8 +180,8 @@ static Optional<Operator> find(String text) {
           .put(MODULO.getFunction(), "%")
           .buildOrThrow();
 
-  /** Lookup an operator by its mangled name, as used within the AST. */
-  static Optional<Operator> findReverse(String op) {
+  /** Lookup an operator by its mangled name (ex: _&&_), as used within the AST. */
+  public static Optional<Operator> findReverse(String op) {
     return Optional.ofNullable(REVERSE_OPERATORS.get(op));
   }
 

parser/src/test/java/dev/cel/parser/OperatorTest.java

@@ -19,6 +19,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 import com.google.testing.junit.testparameterinjector.TestParameters;
 import dev.cel.common.ast.CelExpr;
@@ -50,6 +51,11 @@ public void findReverse_returnsCorrectOperator() {
     assertThat(Operator.findReverse("_+_")).hasValue(Operator.ADD);
   }
 
+  @Test
+  public void findReverse_allOperators(@TestParameter Operator operator) {
+    assertThat(Operator.findReverse(operator.getFunction())).hasValue(operator);
+  }
+
   @Test
   public void findReverseBinaryOperator_returnsEmptyWhenNotFound() {
     assertThat(Operator.findReverseBinaryOperator("+")).isEmpty();

policy/src/main/java/dev/cel/policy/BUILD.bazel

@@ -215,6 +215,8 @@ java_library(
         "//optimizer",
         "//optimizer:optimization_exception",
         "//optimizer:optimizer_builder",
+        "//optimizer/optimizers:common_subexpression_elimination",
+        "//optimizer/optimizers:constant_folding",
         "//validator",
         "//validator:ast_validator",
         "//validator:validator_builder",
@@ -247,7 +249,9 @@ java_library(
         "//common:cel_ast",
         "//common:compiler_common",
         "//common:mutable_ast",
+        "//common/ast",
         "//common/formats:value_string",
+        "//common/navigation:mutable_navigation",
         "//extensions:optional_library",
         "//optimizer:ast_optimizer",
         "//optimizer:mutable_ast",

policy/src/main/java/dev/cel/policy/CelPolicyCompilerImpl.java

@@ -36,6 +36,9 @@
 import dev.cel.optimizer.CelOptimizationException;
 import dev.cel.optimizer.CelOptimizer;
 import dev.cel.optimizer.CelOptimizerFactory;
+import dev.cel.optimizer.optimizers.ConstantFoldingOptimizer;
+import dev.cel.optimizer.optimizers.SubexpressionOptimizer;
+import dev.cel.optimizer.optimizers.SubexpressionOptimizer.SubexpressionOptimizerOptions;
 import dev.cel.policy.CelCompiledRule.CelCompiledMatch;
 import dev.cel.policy.CelCompiledRule.CelCompiledMatch.Result;
 import dev.cel.policy.CelCompiledRule.CelCompiledMatch.Result.Kind;
@@ -98,7 +101,7 @@ public CelCompiledRule compileRule(CelPolicy policy) throws CelPolicyValidationE
   public CelAbstractSyntaxTree compose(CelPolicy policy, CelCompiledRule compiledRule)
       throws CelPolicyValidationException {
     Cel cel = compiledRule.cel();
-    CelOptimizer optimizer =
+    CelOptimizer composingOptimizer =
         CelOptimizerFactory.standardCelOptimizerBuilder(cel)
             .addAstOptimizers(
                 RuleComposer.newInstance(compiledRule, variablesPrefix, iterationLimit))
@@ -110,7 +113,7 @@ public CelAbstractSyntaxTree compose(CelPolicy policy, CelCompiledRule compiledR
       // This is a minimal expression used as a basis of stitching together all the rules into a
       // single graph.
       ast = cel.compile("true").getAst();
-      ast = optimizer.optimize(ast);
+      ast = composingOptimizer.optimize(ast);
     } catch (CelValidationException | CelOptimizationException e) {
       if (e.getCause() instanceof RuleCompositionException) {
         RuleCompositionException re = (RuleCompositionException) e.getCause();
@@ -136,6 +139,24 @@ public CelAbstractSyntaxTree compose(CelPolicy policy, CelCompiledRule compiledR
       throw new CelPolicyValidationException("Unexpected error while composing rules.", e);
     }
 
+    CelOptimizer astOptimizer =
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
+            .addAstOptimizers(
+                ConstantFoldingOptimizer.getInstance(),
+                SubexpressionOptimizer.newInstance(
+                    SubexpressionOptimizerOptions.newBuilder()
+                        .populateMacroCalls(true)
+                        .enableCelBlock(true)
+                        .build()))
+            .build();
+    try {
+      // Optimize the composed graph using const fold and CSE
+      ast = astOptimizer.optimize(ast);
+    } catch (CelOptimizationException e) {
+      throw new CelPolicyValidationException(
+          "Failed to optimize the composed policy. Reason: " + e.getMessage(), e);
+    }
+
     assertAstDepthIsSafe(ast, cel);
 
     return ast;

policy/src/main/java/dev/cel/policy/RuleComposer.java

@@ -19,12 +19,16 @@
 import static java.util.stream.Collectors.toCollection;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import dev.cel.bundle.Cel;
 import dev.cel.common.CelAbstractSyntaxTree;
 import dev.cel.common.CelMutableAst;
 import dev.cel.common.CelValidationException;
+import dev.cel.common.ast.CelExpr.ExprKind.Kind;
 import dev.cel.common.formats.ValueString;
+import dev.cel.common.navigation.CelNavigableMutableAst;
+import dev.cel.common.navigation.CelNavigableMutableExpr;
 import dev.cel.extensions.CelOptionalLibrary.Function;
 import dev.cel.optimizer.AstMutator;
 import dev.cel.optimizer.CelAstOptimizer;
@@ -151,23 +155,37 @@ private RuleOptimizationResult optimizeRule(Cel cel, CelCompiledRule compiledRul
       }
     }
 
-    CelMutableAst result = matchAst;
-    for (CelCompiledVariable variable : Lists.reverse(compiledRule.variables())) {
-      result =
-          astMutator.replaceSubtreeWithNewBindMacro(
-              result,
-              variablePrefix + variable.name(),
-              CelMutableAst.fromCelAst(variable.ast()),
-              result.expr(),
-              result.expr().id(),
-              true);
-    }
+    CelMutableAst result = inlineCompiledVariables(matchAst, compiledRule.variables());
 
     result = astMutator.renumberIdsConsecutively(result);
 
     return RuleOptimizationResult.create(result, isOptionalResult);
   }
 
+  private CelMutableAst inlineCompiledVariables(
+      CelMutableAst ast, List<CelCompiledVariable> compiledVariables) {
+    CelMutableAst mutatedAst = ast;
+    for (CelCompiledVariable compiledVariable : Lists.reverse(compiledVariables)) {
+      String variableName = variablePrefix + compiledVariable.name();
+      ImmutableList<CelNavigableMutableExpr> exprsToReplace =
+          CelNavigableMutableAst.fromAst(mutatedAst)
+              .getRoot()
+              .allNodes()
+              .filter(
+                  node ->
+                      node.expr().getKind().equals(Kind.IDENT)
+                          && node.expr().ident().name().equals(variableName))
+              .collect(toImmutableList());
+
+      for (CelNavigableMutableExpr expr : exprsToReplace) {
+        CelMutableAst variableAst = CelMutableAst.fromCelAst(compiledVariable.ast());
+        mutatedAst = astMutator.replaceSubtree(mutatedAst, variableAst, expr.id());
+      }
+    }
+
+    return mutatedAst;
+  }
+
   static RuleComposer newInstance(
       CelCompiledRule compiledRule, String variablePrefix, int iterationLimit) {
     return new RuleComposer(compiledRule, variablePrefix, iterationLimit);

policy/src/test/java/dev/cel/policy/CelPolicyCompilerImplTest.java

@@ -136,8 +136,8 @@ public void compileYamlPolicy_multilineContainsError_throws(
 
   @Test
   public void compileYamlPolicy_exceedsDefaultAstDepthLimit_throws() throws Exception {
-    String longExpr =
-        "0+1+2+3+4+5+6+7+8+9+10+11+12+13+14+15+16+17+18+19+20+21+22+23+24+25+26+27+28+29+30+31+32+33+34+35+36+37+38+39+40+41+42+43+44+45+46+47+48+49+50";
+    Cel cel = newCel().toCelBuilder().addVar("msg", SimpleType.DYN).build();
+    String longExpr = "msg.b.c.d.e.f";
     String policyContent =
         String.format(
             "name: deeply_nested_ast\n" + "rule:\n" + "  match:\n" + "    - output: %s", longExpr);
@@ -146,11 +146,35 @@ public void compileYamlPolicy_exceedsDefaultAstDepthLimit_throws() throws Except
     CelPolicyValidationException e =
         assertThrows(
             CelPolicyValidationException.class,
-            () -> CelPolicyCompilerFactory.newPolicyCompiler(newCel()).build().compile(policy));
+            () ->
+                CelPolicyCompilerFactory.newPolicyCompiler(cel)
+                    .setAstDepthLimit(5)
+                    .build()
+                    .compile(policy));
+
+    assertThat(e)
+        .hasMessageThat()
+        .isEqualTo("ERROR: <input>:-1:0: AST's depth exceeds the configured limit: 5.");
+  }
 
+  @Test
+  public void compileYamlPolicy_constantFoldingFailure_throwsDuringComposition() throws Exception {
+    String policyContent =
+        "name: ast_with_div_by_zero\n" //
+            + "rule:\n" //
+            + "  match:\n" //
+            + "    - output: 1 / 0";
+    CelPolicy policy = POLICY_PARSER.parse(policyContent);
+
+    CelPolicyValidationException e =
+        assertThrows(
+            CelPolicyValidationException.class,
+            () -> CelPolicyCompilerFactory.newPolicyCompiler(newCel()).build().compile(policy));
     assertThat(e)
         .hasMessageThat()
-        .isEqualTo("ERROR: <input>:-1:0: AST's depth exceeds the configured limit: 50.");
+        .isEqualTo(
+            "Failed to optimize the composed policy. Reason: Constant folding failure. Failed to"
+                + " evaluate subtree due to: evaluation error: / by zero");
   }
 
   @Test