cel-expr · copybara-service · Aug 21, 2024 · Aug 21, 2024
diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java
@@ -551,7 +551,51 @@ private boolean canEliminate(
         && !(navigableExpr.getKind().equals(Kind.LIST)
             && navigableExpr.expr().list().elements().isEmpty())
         && containsEliminableFunctionOnly(navigableExpr)
-        && !ineligibleExprs.contains(navigableExpr.expr());
+        && !ineligibleExprs.contains(navigableExpr.expr())
+        && containsComprehensionIdentInSubexpr(navigableExpr);
+  }
+
+  private boolean containsComprehensionIdentInSubexpr(CelNavigableMutableExpr navExpr) {
+    if (!cseOptions.retainComprehensionStructure()) {
+      return true;
+    }
+
+    if (navExpr.getKind().equals(Kind.COMPREHENSION)) {
+      return true;
+    }
+
+    ImmutableList<CelNavigableMutableExpr> comprehensionIdents =
+        navExpr
+            .allNodes()
+            .filter(
+                node ->
+                    node.getKind().equals(Kind.IDENT)
+                        && (node.expr()
+                                .ident()
+                                .name()
+                                .startsWith(MANGLED_COMPREHENSION_ITER_VAR_PREFIX)
+                            || node.expr()
+                                .ident()
+                                .name()
+                                .startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX)))
+            .collect(toImmutableList());
+
+    if (comprehensionIdents.isEmpty()) {
+      return true;
+    }
+
+    for (CelNavigableMutableExpr ident : comprehensionIdents) {
+      CelNavigableMutableExpr parent = ident.parent().orElse(null);
+      while (parent != null) {
+        if (parent.getKind().equals(Kind.COMPREHENSION)) {
+          return false;
+        }
+
+        parent = parent.parent().orElse(null);
+      }
+    }
+
+    return true;
   }
 
   /**
@@ -628,6 +672,8 @@ public abstract static class SubexpressionOptimizerOptions {
 
     public abstract boolean enableCelBlock();
 
+    public abstract boolean retainComprehensionStructure();
+
     public abstract int subexpressionMaxRecursionDepth();
 
     public abstract ImmutableSet<String> eliminableFunctions();
@@ -684,6 +730,25 @@ public abstract static class Builder {
        */
       public abstract Builder subexpressionMaxRecursionDepth(int value);
 
+      /**
+       * If configured true, SubexpressionOptimizer will not break apart a subexpression containing
+       * a comprehension's iter_var and accu_var without the surrounding comprehension.
+       *
+       * <p>An example expression {@code ['foo'].map(x, [x+x]) + ['foo'].map(x, [x+x, x+x])} is
+       * optimized as (note the common subexpr x+x that leverage the iteration variable):
+       *
+       * <pre>
+       *   Disabled: {@code cel.@block([["foo"], @it0 + @it0], @index0.map(@it0, [@index1])
+       *       + @index0.map(@it0, [@index1, @index1]))}
+       *   Enabled: {@code cel.@block([["foo"]], @index0.map(@it0, [@it0 + @it0])
+       *       + @index0.map(@it0, [@it0 + @it0, @it0 + @it0]))}
+       *  </pre>
+       *
+       * If targeting CEL-Java for the runtime, the recommended setting is to
+       * leave this disabled for maximal optimization efficiency.
+       */
+      public abstract Builder retainComprehensionStructure(boolean value);
+
       abstract ImmutableSet.Builder<String> eliminableFunctionsBuilder();
 
       /**
@@ -718,6 +783,7 @@ public static Builder newBuilder() {
           .iterationLimit(500)
           .populateMacroCalls(false)
           .enableCelBlock(false)
+          .retainComprehensionStructure(false)
           .subexpressionMaxRecursionDepth(0);
     }
 

diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerBaselineTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerBaselineTest.java
@@ -71,6 +71,7 @@ public class SubexpressionOptimizerBaselineTest extends BaselineTestCase {
 
   private static final SubexpressionOptimizerOptions OPTIMIZER_COMMON_OPTIONS =
       SubexpressionOptimizerOptions.newBuilder()
+          .retainComprehensionStructure(false)
           .populateMacroCalls(true)
           .enableCelBlock(true)
           .addEliminableFunctions("pure_custom_func")
@@ -352,6 +353,8 @@ private static CelOptimizer newCseOptimizer(Cel cel, SubexpressionOptimizerOptio
   private enum CseTestOptimizer {
     CASCADED_BINDS(OPTIMIZER_COMMON_OPTIONS.toBuilder().enableCelBlock(false).build()),
     BLOCK_COMMON_SUBEXPR_ONLY(OPTIMIZER_COMMON_OPTIONS),
+    BLOCK_COMMON_SUBEXPR_COMPREHENSION_STRUCTURE_RETAINED(
+        OPTIMIZER_COMMON_OPTIONS.toBuilder().retainComprehensionStructure(true).build()),
     BLOCK_RECURSION_DEPTH_1(
         OPTIMIZER_COMMON_OPTIONS.toBuilder().subexpressionMaxRecursionDepth(1).build()),
     BLOCK_RECURSION_DEPTH_2(

diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java
@@ -201,6 +201,26 @@ public void cse_withCelBlock_noop(@TestParameter CseNoOpTestCase testCase) throw
     assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(testCase.source);
   }
 
+  @Test
+  public void cse_withComprehensionStructureRetained() throws Exception {
+    CelAbstractSyntaxTree ast =
+        CEL.compile("['foo'].map(x, [x+x]) + ['foo'].map(x, [x+x, x+x])").getAst();
+    CelOptimizer celOptimizer =
+        newCseOptimizer(
+            SubexpressionOptimizerOptions.newBuilder()
+                .populateMacroCalls(true)
+                .enableCelBlock(true)
+                .retainComprehensionStructure(true)
+                .build());
+
+    CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst))
+        .isEqualTo(
+            "cel.@block([[\"foo\"]], @index0.map(@it:0:0, [@it:0:0 + @it:0:0]) +"
+                + " @index0.map(@it:0:0, [@it:0:0 + @it:0:0, @it:0:0 + @it:0:0]))");
+  }
+
   @Test
   public void cse_applyConstFoldingAfter() throws Exception {
     CelAbstractSyntaxTree ast =