Allow specifying a set of optimizers to run to the policy compiler

PiperOrigin-RevId: 885225062

Author: l46kok

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

@@ -138,6 +138,7 @@ java_library(
     ],
     deps = [
         ":compiler",
+        "//optimizer:ast_optimizer",
         "@maven//:com_google_errorprone_error_prone_annotations",
     ],
 )
@@ -214,6 +215,7 @@ java_library(
         "//common/types",
         "//common/types:type_providers",
         "//optimizer",
+        "//optimizer:ast_optimizer",
         "//optimizer:optimization_exception",
         "//optimizer:optimizer_builder",
         "//optimizer/optimizers:common_subexpression_elimination",

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

@@ -16,6 +16,8 @@
 
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
+import dev.cel.optimizer.CelAstOptimizer;
+import java.util.List;
 
 /** Interface for building an instance of {@link CelPolicyCompiler} */
 public interface CelPolicyCompilerBuilder {
@@ -38,6 +40,10 @@ public interface CelPolicyCompilerBuilder {
   @CanIgnoreReturnValue
   CelPolicyCompilerBuilder setAstDepthLimit(int iterationLimit);
 
+  /** Configures the policy compiler to run the provided optimizers on compiled policies. */
+  @CanIgnoreReturnValue
+  CelPolicyCompilerBuilder setOptimizers(List<CelAstOptimizer> optimizers);
+
   @CheckReturnValue
   CelPolicyCompiler build();
 }

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

@@ -33,6 +33,7 @@
 import dev.cel.common.formats.ValueString;
 import dev.cel.common.types.CelType;
 import dev.cel.common.types.SimpleType;
+import dev.cel.optimizer.CelAstOptimizer;
 import dev.cel.optimizer.CelOptimizationException;
 import dev.cel.optimizer.CelOptimizer;
 import dev.cel.optimizer.CelOptimizerFactory;
@@ -63,6 +64,7 @@ final class CelPolicyCompilerImpl implements CelPolicyCompiler {
   private final Cel cel;
   private final String variablesPrefix;
   private final int iterationLimit;
+  private final ImmutableList<CelAstOptimizer> optimizers;
   private final Optional<CelAstValidator> astDepthValidator;
 
   @Override
@@ -140,19 +142,7 @@ public CelAbstractSyntaxTree compose(CelPolicy policy, CelCompiledRule compiledR
     }
 
     CelOptimizer astOptimizer =
-        CelOptimizerFactory.standardCelOptimizerBuilder(cel)
-            .addAstOptimizers(
-                ConstantFoldingOptimizer.getInstance(),
-                SubexpressionOptimizer.newInstance(
-                    SubexpressionOptimizerOptions.newBuilder()
-                        // "record" is used for recording subexpression results via
-                        // BlueprintLateFunctionBinding. Safely eliminable, since repeated
-                        // invocation does not change the intermediate results.
-                        .addEliminableFunctions("record")
-                        .populateMacroCalls(true)
-                        .enableCelBlock(true)
-                        .build()))
-            .build();
+        CelOptimizerFactory.standardCelOptimizerBuilder(cel).addAstOptimizers(optimizers).build();
     try {
       // Optimize the composed graph using const fold and CSE
       ast = astOptimizer.optimize(ast);
@@ -339,6 +329,7 @@ static final class Builder implements CelPolicyCompilerBuilder {
     private final Cel cel;
     private String variablesPrefix;
     private int iterationLimit;
+    private ImmutableList<CelAstOptimizer> optimizers;
     private Optional<CelAstValidator> astDepthLimitValidator;
 
     private Builder(Cel cel) {
@@ -362,7 +353,7 @@ public Builder setIterationLimit(int iterationLimit) {
 
     @Override
     @CanIgnoreReturnValue
-    public CelPolicyCompilerBuilder setAstDepthLimit(int astDepthLimit) {
+    public Builder setAstDepthLimit(int astDepthLimit) {
       if (astDepthLimit < 0) {
         astDepthLimitValidator = Optional.empty();
       } else {
@@ -371,27 +362,40 @@ public CelPolicyCompilerBuilder setAstDepthLimit(int astDepthLimit) {
       return this;
     }
 
+    @Override
+    public Builder setOptimizers(List<CelAstOptimizer> optimizers) {
+      this.optimizers = ImmutableList.copyOf(optimizers);
+      return this;
+    }
+
     @Override
     public CelPolicyCompiler build() {
       return new CelPolicyCompilerImpl(
-          cel, this.variablesPrefix, this.iterationLimit, astDepthLimitValidator);
+          cel, this.variablesPrefix, this.iterationLimit, this.optimizers, astDepthLimitValidator);
     }
   }
 
   static Builder newBuilder(Cel cel) {
     return new Builder(cel)
         .setVariablesPrefix(DEFAULT_VARIABLE_PREFIX)
-        .setIterationLimit(DEFAULT_ITERATION_LIMIT);
+        .setIterationLimit(DEFAULT_ITERATION_LIMIT)
+        .setOptimizers(
+            ImmutableList.of(
+                ConstantFoldingOptimizer.getInstance(),
+                SubexpressionOptimizer.newInstance(
+                    SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build())));
   }
 
   private CelPolicyCompilerImpl(
       Cel cel,
       String variablesPrefix,
       int iterationLimit,
+      ImmutableList<CelAstOptimizer> optimizers,
       Optional<CelAstValidator> astDepthValidator) {
     this.cel = checkNotNull(cel);
     this.variablesPrefix = checkNotNull(variablesPrefix);
     this.iterationLimit = iterationLimit;
+    this.optimizers = optimizers;
     this.astDepthValidator = astDepthValidator;
   }
 }