Remove cel.bind option from SubexpressionOptimizer

PiperOrigin-RevId: 797496424

Author: l46kok

optimizer/src/main/java/dev/cel/optimizer/AstMutator.java

@@ -910,21 +910,32 @@ private static void unwrapListArgumentsInMacroCallExpr(
 
     CelMutableExpr loopStepExpr = comprehension.loopStep();
     List<CelMutableExpr> loopStepArgs = loopStepExpr.call().args();
-    if (loopStepArgs.size() != 2) {
+    if (loopStepArgs.size() != 2 && loopStepArgs.size() != 3) {
       throw new IllegalArgumentException(
           String.format(
-              "Expected exactly 2 arguments but got %d instead on expr id: %d",
+              "Expected exactly 2 or 3 arguments but got %d instead on expr id: %d",
               loopStepArgs.size(), loopStepExpr.id()));
     }
 
     CelMutableCall existingMacroCall = newMacroCallExpr.call();
     CelMutableCall newMacroCall =
         existingMacroCall.target().isPresent()
-            ? CelMutableCall.create(existingMacroCall.target().get(), existingMacroCall.function())
+            ? CelMutableCall.create(existingMacroCall.target().get(),
+            existingMacroCall.function())
             : CelMutableCall.create(existingMacroCall.function());
     newMacroCall.addArgs(
         existingMacroCall.args().get(0)); // iter_var is first argument of the call by convention
-    newMacroCall.addArgs(loopStepArgs.get(1).list().elements());
+
+    CelMutableList extraneousList = null;
+    if (loopStepArgs.size() == 2) {
+      extraneousList = loopStepArgs.get(1).list();
+    } else {
+      newMacroCall.addArgs(loopStepArgs.get(0));
+      // For map(x,y,z), z is wrapped in a _+_(@result, [z])
+      extraneousList = loopStepArgs.get(1).call().args().get(1).list();
+    }
+
+    newMacroCall.addArgs(extraneousList.elements());
 
     newMacroCallExpr.setCall(newMacroCall);
   }

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.

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

@@ -57,10 +57,8 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 
 /**
@@ -121,8 +119,7 @@ public static SubexpressionOptimizer newInstance(SubexpressionOptimizerOptions c
 
   @Override
   public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) {
-    OptimizationResult result =
-        cseOptions.enableCelBlock() ? optimizeUsingCelBlock(ast, cel) : optimizeUsingCelBind(ast);
+    OptimizationResult result = optimizeUsingCelBlock(ast, cel);
 
     verifyOptimizedAstCorrectness(result.optimizedAst());
 
@@ -184,7 +181,7 @@ private OptimizationResult optimizeUsingCelBlock(CelAbstractSyntaxTree ast, Cel
 
     if (iterCount == 0) {
       // No modification has been made.
-      return OptimizationResult.create(astToModify.toParsedAst());
+      return OptimizationResult.create(ast);
     }
 
     ImmutableList.Builder<CelVarDecl> newVarDecls = ImmutableList.builder();
@@ -330,123 +327,8 @@ private static ImmutableList<CelVarDecl> newBlockIndexVariableDeclarations(
     return varDeclBuilder.build();
   }
 
-  private OptimizationResult optimizeUsingCelBind(CelAbstractSyntaxTree ast) {
-    CelMutableAst astToModify = CelMutableAst.fromCelAst(ast);
-    if (!cseOptions.populateMacroCalls()) {
-      astToModify.source().clearMacroCalls();
-    }
-
-    astToModify =
-        astMutator
-            .mangleComprehensionIdentifierNames(
-                astToModify,
-                MANGLED_COMPREHENSION_ITER_VAR_PREFIX,
-                MANGLED_COMPREHENSION_ACCU_VAR_PREFIX,
-                /* incrementSerially= */ true)
-            .mutableAst();
-    CelMutableSource sourceToModify = astToModify.source();
-
-    int bindIdentifierIndex = 0;
-    int iterCount;
-    for (iterCount = 0; iterCount < cseOptions.iterationLimit(); iterCount++) {
-      CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(astToModify);
-      List<CelMutableExpr> cseCandidates = getCseCandidates(navAst);
-      if (cseCandidates.isEmpty()) {
-        break;
-      }
-
-      String bindIdentifier = BIND_IDENTIFIER_PREFIX + bindIdentifierIndex;
-      bindIdentifierIndex++;
-
-      // Replace all CSE candidates with new bind identifier
-      for (CelMutableExpr cseCandidate : cseCandidates) {
-        iterCount++;
-
-        astToModify =
-            astMutator.replaceSubtree(
-                astToModify, CelMutableExpr.ofIdent(bindIdentifier), cseCandidate.id());
-      }
-
-      // Find LCA to insert the new cel.bind macro into.
-      CelNavigableMutableExpr lca = getLca(navAst, bindIdentifier);
-
-      // Insert the new bind call
-      CelMutableExpr subexpressionToBind = cseCandidates.get(0);
-      // Re-add the macro source for bind identifiers that might have been lost from previous
-      // iteration of CSE
-      astToModify.source().addAllMacroCalls(sourceToModify.getMacroCalls());
-      astToModify =
-          astMutator.replaceSubtreeWithNewBindMacro(
-              astToModify,
-              bindIdentifier,
-              subexpressionToBind,
-              lca.expr(),
-              lca.id(),
-              cseOptions.populateMacroCalls());
-
-      // Retain the existing macro calls in case if the bind identifiers are replacing a subtree
-      // that contains a comprehension.
-      sourceToModify = astToModify.source();
-    }
-
-    if (iterCount >= cseOptions.iterationLimit()) {
-      throw new IllegalStateException("Max iteration count reached.");
-    }
-
-    if (iterCount == 0) {
-      // No modification has been made.
-      return OptimizationResult.create(astToModify.toParsedAst());
-    }
-
-    astToModify = astMutator.renumberIdsConsecutively(astToModify);
-
-    return OptimizationResult.create(astToModify.toParsedAst());
-  }
-
-  private static CelNavigableMutableExpr getLca(
-      CelNavigableMutableAst navAst, String boundIdentifier) {
-    CelNavigableMutableExpr root = navAst.getRoot();
-    ImmutableList<CelNavigableMutableExpr> allNodesWithIdentifier =
-        root.allNodes()
-            .filter(
-                node ->
-                    node.getKind().equals(Kind.IDENT)
-                        && node.expr().ident().name().equals(boundIdentifier))
-            .collect(toImmutableList());
-
-    if (allNodesWithIdentifier.size() < 2) {
-      throw new IllegalStateException("Expected at least 2 bound identifiers to be present.");
-    }
-
-    CelNavigableMutableExpr lca = root;
-    long lcaAncestorCount = 0;
-    HashMap<Long, Long> ancestors = new HashMap<>();
-    for (CelNavigableMutableExpr navigableExpr : allNodesWithIdentifier) {
-      Optional<CelNavigableMutableExpr> maybeParent = Optional.of(navigableExpr);
-      while (maybeParent.isPresent()) {
-        CelNavigableMutableExpr parent = maybeParent.get();
-        if (!ancestors.containsKey(parent.id())) {
-          ancestors.put(parent.id(), 1L);
-          continue;
-        }
-
-        long ancestorCount = ancestors.get(parent.id());
-        if (lcaAncestorCount < ancestorCount
-            || (lcaAncestorCount == ancestorCount && lca.depth() < parent.depth())) {
-          lca = parent;
-          lcaAncestorCount = ancestorCount;
-        }
-
-        ancestors.put(parent.id(), ancestorCount + 1);
-        maybeParent = parent.parent();
-      }
-    }
-
-    return lca;
-  }
-
   private List<CelMutableExpr> getCseCandidates(CelNavigableMutableAst navAst) {
-    if (cseOptions.enableCelBlock() && cseOptions.subexpressionMaxRecursionDepth() > 0) {
+    if (cseOptions.subexpressionMaxRecursionDepth() > 0) {
       return getCseCandidatesWithRecursionDepth(
           navAst, cseOptions.subexpressionMaxRecursionDepth());
     } else {
@@ -689,11 +571,9 @@ public abstract static class Builder {
       public abstract Builder populateMacroCalls(boolean value);
 
       /**
-       * Rewrites the optimized AST using cel.@block call instead of cascaded cel.bind macros, aimed
-       * to produce a more compact AST. {@link CelSource.Extension} field will be populated in the
-       * AST to inform that special runtime support is required to evaluate the optimized
-       * expression.
+       * @deprecated This option is a no-op. cel.@block is always enabled.
        */
+      @Deprecated
       public abstract Builder enableCelBlock(boolean value);
 
       /**
@@ -708,9 +588,8 @@ public abstract static class Builder {
        * <p>Note that expressions containing no common subexpressions may become a candidate for
        * extraction to satisfy the max depth requirement.
        *
-       * <p>This is a no-op if {@link #enableCelBlock} is set to false, the configured value is less
-       * than 1, or no subexpression needs to be extracted because the entire expression is already
-       * under the designated limit.
+       * <p>This is a no-op if the configured value is less than 1, or no subexpression needs to be
+       * extracted because the entire expression is already under the designated limit.
        *
        * <p>Examples:
        *
@@ -756,8 +635,8 @@ public Builder addEliminableFunctions(String... functions) {
     public static Builder newBuilder() {
       return new AutoValue_SubexpressionOptimizer_SubexpressionOptimizerOptions.Builder()
           .iterationLimit(500)
+          .enableCelBlock(true)
           .populateMacroCalls(false)
-          .enableCelBlock(false)
           .subexpressionMaxRecursionDepth(0);
     }
 

optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java

@@ -581,6 +581,34 @@ public void replaceSubtree_replaceExtraneousListCreatedByMacro_unparseSuccess()
         .containsExactly(2L);
   }
 
+  @Test
+  @SuppressWarnings("unchecked") // Test only
+  public void replaceSubtree_replaceExtraneousListCreatedByThreeArgMacro_unparseSuccess()
+      throws Exception {
+    CelAbstractSyntaxTree ast = CEL.compile("[1].map(x, true, 1)").getAst();
+    CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast);
+    CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast);
+
+    // These two mutation are equivalent.
+    CelAbstractSyntaxTree mutatedAstWithList =
+        AST_MUTATOR
+            .replaceSubtree(
+                mutableAst,
+                CelMutableExpr.ofList(
+                    CelMutableList.create(CelMutableExpr.ofConstant(CelConstant.ofValue(2L)))),
+                10L)
+            .toParsedAst();
+    CelAbstractSyntaxTree mutatedAstWithConstant =
+        AST_MUTATOR
+            .replaceSubtree(mutableAst2, CelMutableExpr.ofConstant(CelConstant.ofValue(2L)), 6L)
+            .toParsedAst();
+
+    assertThat(CEL_UNPARSER.unparse(mutatedAstWithList)).isEqualTo("[1].map(x, true, 2)");
+    assertThat(CEL_UNPARSER.unparse(mutatedAstWithConstant)).isEqualTo("[1].map(x, true, 2)");
+    assertThat((List<Long>) CEL.createProgram(CEL.check(mutatedAstWithList).getAst()).eval())
+        .containsExactly(2L);
+  }
+
   @Test
   public void globalCallExpr_replaceRoot() throws Exception {
     // Tree shape (brackets are expr IDs):

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

@@ -73,7 +73,6 @@ public class SubexpressionOptimizerBaselineTest extends BaselineTestCase {
   private static final SubexpressionOptimizerOptions OPTIMIZER_COMMON_OPTIONS =
       SubexpressionOptimizerOptions.newBuilder()
           .populateMacroCalls(true)
-          .enableCelBlock(true)
           .addEliminableFunctions("pure_custom_func")
           .build();
 
@@ -187,67 +186,11 @@ public void subexpression_ast(@TestParameter CseTestOptimizer cseTestOptimizer)
     }
   }
 
-  @Test
-  public void populateMacroCallsDisabled_macroMapUnpopulated(@TestParameter CseTestCase testCase)
-      throws Exception {
-    skipBaselineVerification();
-    Cel cel = newCelBuilder().build();
-    CelOptimizer celOptimizerWithBinds =
-        newCseOptimizer(
-            cel,
-            SubexpressionOptimizerOptions.newBuilder()
-                .populateMacroCalls(false)
-                .enableCelBlock(false)
-                .build());
-    CelOptimizer celOptimizerWithBlocks =
-        newCseOptimizer(
-            cel,
-            SubexpressionOptimizerOptions.newBuilder()
-                .populateMacroCalls(false)
-                .enableCelBlock(true)
-                .build());
-    CelOptimizer celOptimizerWithFlattenedBlocks =
-        newCseOptimizer(
-            cel,
-            SubexpressionOptimizerOptions.newBuilder()
-                .populateMacroCalls(false)
-                .enableCelBlock(true)
-                .subexpressionMaxRecursionDepth(1)
-                .build());
-    CelAbstractSyntaxTree originalAst = cel.compile(testCase.source).getAst();
-
-    CelAbstractSyntaxTree astOptimizedWithBinds = celOptimizerWithBinds.optimize(originalAst);
-    CelAbstractSyntaxTree astOptimizedWithBlocks = celOptimizerWithBlocks.optimize(originalAst);
-    CelAbstractSyntaxTree astOptimizedWithFlattenedBlocks =
-        celOptimizerWithFlattenedBlocks.optimize(originalAst);
-
-    assertThat(astOptimizedWithBinds.getSource().getMacroCalls()).isEmpty();
-    assertThat(astOptimizedWithBlocks.getSource().getMacroCalls()).isEmpty();
-    assertThat(astOptimizedWithFlattenedBlocks.getSource().getMacroCalls()).isEmpty();
-  }
-
-  @Test
-  public void large_expressions_bind_cascaded() throws Exception {
-    CelOptimizer celOptimizer =
-        newCseOptimizer(
-            CEL,
-            SubexpressionOptimizerOptions.newBuilder()
-                .populateMacroCalls(true)
-                .enableCelBlock(false)
-                .build());
-
-    runLargeTestCases(celOptimizer);
-  }
-
   @Test
   public void large_expressions_block_common_subexpr() throws Exception {
     CelOptimizer celOptimizer =
         newCseOptimizer(
-            CEL,
-            SubexpressionOptimizerOptions.newBuilder()
-                .populateMacroCalls(true)
-                .enableCelBlock(true)
-                .build());
+            CEL, SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build());
 
     runLargeTestCases(celOptimizer);
   }
@@ -259,7 +202,6 @@ public void large_expressions_block_recursion_depth_1() throws Exception {
             CEL,
             SubexpressionOptimizerOptions.newBuilder()
                 .populateMacroCalls(true)
-                .enableCelBlock(true)
                 .subexpressionMaxRecursionDepth(1)
                 .build());
 
@@ -273,7 +215,6 @@ public void large_expressions_block_recursion_depth_2() throws Exception {
             CEL,
             SubexpressionOptimizerOptions.newBuilder()
                 .populateMacroCalls(true)
-                .enableCelBlock(true)
                 .subexpressionMaxRecursionDepth(2)
                 .build());
 
@@ -287,7 +228,6 @@ public void large_expressions_block_recursion_depth_3() throws Exception {
             CEL,
             SubexpressionOptimizerOptions.newBuilder()
                 .populateMacroCalls(true)
-                .enableCelBlock(true)
                 .subexpressionMaxRecursionDepth(3)
                 .build());
 
@@ -351,7 +291,6 @@ private static CelOptimizer newCseOptimizer(Cel cel, SubexpressionOptimizerOptio
 
   @SuppressWarnings("Immutable") // Test only
   private enum CseTestOptimizer {
-    CASCADED_BINDS(OPTIMIZER_COMMON_OPTIONS.toBuilder().enableCelBlock(false).build()),
     BLOCK_COMMON_SUBEXPR_ONLY(OPTIMIZER_COMMON_OPTIONS),
     BLOCK_RECURSION_DEPTH_1(
         OPTIMIZER_COMMON_OPTIONS.toBuilder().subexpressionMaxRecursionDepth(1).build()),