Add deny-list mode to SubexpressionOptimizer and enable it in CelPolicyCompilerTool

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

@@ -562,7 +562,10 @@ private boolean containsEliminableFunctionOnly(CelNavigableMutableExpr navigable
         .allMatch(
             node -> {
               if (node.getKind().equals(Kind.CALL)) {
-                return cseEliminableFunctions.contains(node.expr().call().function());
+                String function = node.expr().call().function();
+                return !cseOptions.nonEliminableFunctions().contains(function)
+                    && (cseOptions.eliminateAllFunctions()
+                        || cseEliminableFunctions.contains(function));
               }
 
               return true;
@@ -722,6 +725,10 @@ public abstract static class SubexpressionOptimizerOptions {
 
     public abstract ImmutableSet<String> eliminableFunctions();
 
+    public abstract boolean eliminateAllFunctions();
+
+    public abstract ImmutableSet<String> nonEliminableFunctions();
+
     /** Builder for configuring the {@link SubexpressionOptimizerOptions}. */
     @AutoValue.Builder
     public abstract static class Builder {
@@ -751,20 +758,44 @@ public abstract static class Builder {
        * elimination. By default, standard functions are eliminable.
        *
        * <p>Note that the implementation of custom functions must be free of side effects.
+       *
+       * @deprecated Use {@code eliminateAllFunctions(true)} and {@link #addNonEliminableFunctions}
+       *     to exclude functions with side effects instead.
        */
+      @Deprecated
       @CanIgnoreReturnValue
       public Builder addEliminableFunctions(Iterable<String> functions) {
         checkNotNull(functions);
         this.eliminableFunctionsBuilder().addAll(functions);
         return this;
       }
 
-      /** See {@link #addEliminableFunctions(Iterable)}. */
+      /**
+       * @deprecated See {@link #addEliminableFunctions(Iterable)}.
+       */
+      @Deprecated
       @CanIgnoreReturnValue
       public Builder addEliminableFunctions(String... functions) {
         return addEliminableFunctions(Arrays.asList(functions));
       }
 
+      /**
+       * Treats all functions, including custom ones, as candidates for common subexpression
+       * elimination. Functions with side effects must be excluded via {@link
+       * #addNonEliminableFunctions}.
+       */
+      public abstract Builder eliminateAllFunctions(boolean value);
+
+      abstract ImmutableSet.Builder<String> nonEliminableFunctionsBuilder();
+
+      /** Excludes functions (e.g. those with side effects) from elimination. */
+      @CanIgnoreReturnValue
+      public Builder addNonEliminableFunctions(Iterable<String> functions) {
+        checkNotNull(functions);
+        this.nonEliminableFunctionsBuilder().addAll(functions);
+        return this;
+      }
+
       public abstract SubexpressionOptimizerOptions build();
 
       Builder() {}
@@ -775,7 +806,8 @@ public static Builder newBuilder() {
       return new AutoValue_SubexpressionOptimizer_SubexpressionOptimizerOptions.Builder()
           .iterationLimit(500)
           .enableCelBlock(true)
-          .populateMacroCalls(false);
+          .populateMacroCalls(false)
+          .eliminateAllFunctions(false);
     }
 
     SubexpressionOptimizerOptions() {}

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

@@ -115,7 +115,8 @@ private static Cel setupCelEnv(CelBuilder celBuilder) {
   private static final SubexpressionOptimizerOptions OPTIMIZER_COMMON_OPTIONS =
       SubexpressionOptimizerOptions.newBuilder()
           .populateMacroCalls(true)
-          .addEliminableFunctions("pure_custom_func")
+          .eliminateAllFunctions(true)
+          .addNonEliminableFunctions(ImmutableList.of("non_pure_custom_func"))
           .build();
 
   @Before

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

@@ -200,6 +200,18 @@ public void cse_withUnknownAttributes() throws Exception {
     assertThat(result).isInstanceOf(CelUnknownSet.class);
   }
 
+  @Test
+  public void cse_defaultOptions_macroCallsNotPopulated() throws Exception {
+    CelAbstractSyntaxTree ast =
+        cel.compile("[1].exists(x, x > 0) == [1].exists(x, x > 0)").getAst();
+
+    CelAbstractSyntaxTree optimizedAst =
+        newCseOptimizer(SubexpressionOptimizerOptions.newBuilder().build()).optimize(ast);
+
+    assertThat(ast.getSource().getMacroCalls()).isNotEmpty();
+    assertThat(optimizedAst.getSource().getMacroCalls()).isEmpty();
+  }
+
   private enum CseNoOpTestCase {
     // Nothing to optimize
     NO_COMMON_SUBEXPR("size(\"hello\")"),
@@ -252,6 +264,39 @@ public void cse_withCelBlock_noop(@TestParameter CseNoOpTestCase testCase) throw
     assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(testCase.source);
   }
 
+  @Test
+  public void cse_eliminateAllFunctions_customFunctionEliminated() throws Exception {
+    CelAbstractSyntaxTree ast =
+        cel.compile("non_pure_custom_func(1) + non_pure_custom_func(1)").getAst();
+
+    CelAbstractSyntaxTree optimizedAst =
+        newCseOptimizer(
+                SubexpressionOptimizerOptions.newBuilder().eliminateAllFunctions(true).build())
+            .optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst))
+        .isEqualTo("cel.@block([non_pure_custom_func(1)], @index0 + @index0)");
+  }
+
+  @Test
+  public void cse_eliminateAllFunctions_nonEliminableFunctionRetained() throws Exception {
+    CelAbstractSyntaxTree ast =
+        cel.compile("non_pure_custom_func(size('a')) + non_pure_custom_func(size('a'))").getAst();
+
+    CelAbstractSyntaxTree optimizedAst =
+        newCseOptimizer(
+                SubexpressionOptimizerOptions.newBuilder()
+                    .eliminateAllFunctions(true)
+                    .addNonEliminableFunctions(ImmutableList.of("non_pure_custom_func"))
+                    .build())
+            .optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst))
+        .isEqualTo(
+            "cel.@block([size(\"a\")], non_pure_custom_func(@index0) +"
+                + " non_pure_custom_func(@index0))");
+  }
+
   @Test
   public void cse_withComprehensionStructureRetained() throws Exception {
     CelAbstractSyntaxTree ast =

policy/src/main/java/dev/cel/policy/tools/CelPolicyCompilerTool.java

@@ -133,6 +133,14 @@ public final class CelPolicyCompilerTool implements Callable<Integer> {
       description = "Maximum iteration limit for composing and optimizing the policy")
   private int iterationLimit = 1000;
 
+  @Option(
+      names = {"--non_eliminable_functions"},
+      split = ",",
+      description =
+          "Comma-separated functions with side effects to exclude from common subexpression"
+              + " elimination")
+  private String[] nonEliminableFunctions = {};
+
   private static final CelOptions CEL_OPTIONS =
       CelOptions.current()
           .populateMacroCalls(true)
@@ -233,6 +241,9 @@ public Integer call() {
                       SubexpressionOptimizerOptions.newBuilder()
                           .iterationLimit(iterationLimit)
                           .populateMacroCalls(true)
+                          // CEL functions are assumed side-effect free unless opted out.
+                          .eliminateAllFunctions(true)
+                          .addNonEliminableFunctions(ImmutableList.copyOf(nonEliminableFunctions))
                           .build()));
       if (optimizeFieldSelection) {
         optimizersBuilder.add(

policy/src/test/java/dev/cel/policy/tools/BUILD.bazel

@@ -37,6 +37,7 @@ java_library(
         "//common:cel_source",
         "//common:proto_ast",
         "//extensions:optional_library",
+        "//parser:unparser",
         "//policy/tools",
         "//runtime",
         "@bazel_tools//tools/java/runfiles",

policy/src/test/java/dev/cel/policy/tools/CelPolicyCompilerToolTest.java

@@ -18,6 +18,7 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import dev.cel.expr.CheckedExpr;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.io.Files;
 import com.google.devtools.build.runfiles.AutoBazelRepository;
@@ -29,6 +30,8 @@
 import dev.cel.common.CelSource;
 import dev.cel.expr.conformance.proto3.TestAllTypes;
 import dev.cel.extensions.CelOptionalLibrary;
+import dev.cel.parser.CelUnparser;
+import dev.cel.parser.CelUnparserFactory;
 import dev.cel.runtime.CelRuntime;
 import dev.cel.runtime.CelRuntimeFactory;
 import java.io.ByteArrayOutputStream;
@@ -50,6 +53,8 @@
 @AutoBazelRepository
 public final class CelPolicyCompilerToolTest {
 
+  private static final CelUnparser CEL_UNPARSER = CelUnparserFactory.newUnparser();
+
   @Rule public TemporaryFolder tempFolder = new TemporaryFolder();
 
   private Runfiles runfiles;
@@ -327,6 +332,27 @@ public void compile_withIterationLimitReached_returnsError() throws Exception {
     assertThat(stdErr).contains("Reason: Unexpected error while composing rules.");
   }
 
+  @Test
+  public void compile_customFunction_eliminatedByDefault() throws Exception {
+    CelAbstractSyntaxTree ast = compileCustomFunctionPolicy();
+
+    assertThat(CEL_UNPARSER.unparse(ast))
+        .isEqualTo(
+            "cel.@block([expensive(\"x\")], (@index0 > 1 && @index0 < 10) ? optional.of(true) :"
+                + " optional.none())");
+  }
+
+  @Test
+  public void compile_withNonEliminableFunctions_retainsCalls() throws Exception {
+    CelAbstractSyntaxTree ast =
+        compileCustomFunctionPolicy("--non_eliminable_functions", "expensive,unused");
+
+    assertThat(CEL_UNPARSER.unparse(ast))
+        .isEqualTo(
+            "(expensive(\"x\") > 1 && expensive(\"x\") < 10) ? optional.of(true) :"
+                + " optional.none()");
+  }
+
   @Test
   public void compile_withOptimizeFieldSelection_rewritesSelectAndEvaluates() throws Exception {
     String configRlocation =
@@ -923,6 +949,47 @@ private String createFile(String fileName, String content) throws IOException {
     return file.getAbsolutePath();
   }
 
+  private CelAbstractSyntaxTree compileCustomFunctionPolicy(String... extraArgs) throws Exception {
+    String configPath =
+        createFile(
+            "config.yaml",
+            "name: test-env\n"
+                + "functions:\n"
+                + "  - name: expensive\n"
+                + "    overloads:\n"
+                + "      - id: expensive_string\n"
+                + "        args:\n"
+                + "          - type_name: string\n"
+                + "        return:\n"
+                + "          type_name: int\n");
+    String policyPath =
+        createFile(
+            "policy.yaml",
+            "name: p\n"
+                + "rule:\n"
+                + "  variables:\n"
+                + "    - name: score\n"
+                + "      expression: expensive('x')\n"
+                + "  match:\n"
+                + "    - condition: variables.score > 1 && variables.score < 10\n"
+                + "      output: 'true'\n");
+    File outputFile = tempFolder.newFile("output.binarypb");
+    ImmutableList<String> args =
+        ImmutableList.<String>builder()
+            .add("--policy", policyPath, "--config", configPath)
+            .add("--output", outputFile.getAbsolutePath())
+            .add(extraArgs)
+            .build();
+
+    int exitCode =
+        new CommandLine(new CelPolicyCompilerTool()).execute(args.toArray(new String[0]));
+
+    assertThat(exitCode).isEqualTo(0);
+    CheckedExpr checkedExpr =
+        CheckedExpr.parseFrom(Files.toByteArray(outputFile), ExtensionRegistry.getEmptyRegistry());
+    return CelProtoAbstractSyntaxTree.fromCheckedExpr(checkedExpr).getAst();
+  }
+
   private static String executeExpectingError(String... args) {
     StringWriter out = new StringWriter();
     PrintWriter pw = new PrintWriter(out);

policy/tools/compile_cel_policy.bzl

@@ -28,6 +28,7 @@ def compile_cel_policy(
         optimize_field_selection = False,
         simple_variables = False,
         iteration_limit = None,
+        non_eliminable_functions = [],
         visibility = None):
     """Compiles a CEL policy into a CheckedExpr binarypb or textpb with optional select optimization.
 
@@ -54,6 +55,8 @@ def compile_cel_policy(
         `- var_name: expr` instead of `- name: var_name` / `expression: expr`) in the policy (default False)
       iteration_limit: (optional) int maximum iteration limit for composing and optimizing the policy
         (default None, which uses the compiler default of 1000)
+      non_eliminable_functions: (optional) list of str function names with side effects to exclude
+        from common subexpression elimination. All other functions are assumed side-effect free.
       visibility: (optional) visibility to use on the genrule macro (default None)
     """
     if output_format not in ("binarypb", "textpb", "textproto"):
@@ -68,11 +71,11 @@ def compile_cel_policy(
     args = []
     srcs = [policy, config]
 
-    args.append("--policy=$(location %s)" % policy)
-    args.append("--config=$(location %s)" % config)
+    args.append("--policy=$(execpath %s)" % policy)
+    args.append("--config=$(execpath %s)" % config)
 
     if base_config != None:
-        args.append("--base_config=$(location %s)" % base_config)
+        args.append("--base_config=$(execpath %s)" % base_config)
         srcs.append(base_config)
 
     if len(proto_srcs) > 0 and file_descriptor_set != None:
@@ -87,10 +90,10 @@ def compile_cel_policy(
         file_descriptor_set = transitive_descriptor_set_name
 
     if file_descriptor_set != None:
-        args.append("--file_descriptor_set=$(location %s)" % file_descriptor_set)
+        args.append("--file_descriptor_set=$(execpath %s)" % file_descriptor_set)
         srcs.append(file_descriptor_set)
 
-    args.append("--output=$(location %s)" % output)
+    args.append("--output=$(execpath %s)" % output)
     args.append("--output_format=" + output_format)
     args.append("--output_version=" + output_version)
 
@@ -103,8 +106,11 @@ def compile_cel_policy(
     if iteration_limit != None:
         args.append("--iteration_limit=%d" % iteration_limit)
 
+    if non_eliminable_functions:
+        args.append("--non_eliminable_functions=" + ",".join(non_eliminable_functions))
+
     cmd = (
-        "$(location //policy/tools:cel_policy_compiler_tool) " +
+        "$(execpath //policy/tools:cel_policy_compiler_tool) " +
         " ".join(args)
     )