Hoist loop-invariant subexpressions in SubexpressionOptimizer

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

@@ -26,11 +26,13 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import com.google.common.base.Verify;
+import com.google.common.collect.HashMultiset;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Multimaps;
+import com.google.common.collect.Multiset;
 import com.google.common.collect.Sets;
 import com.google.common.collect.Streams;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -414,15 +416,40 @@ private ImmutableList<CelMutableExpr> getCseCandidates(
         navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).collect(toImmutableList());
     ImmutableListMultimap<Integer, CelNavigableMutableExpr> nodesByHeight =
         Multimaps.index(allNodes, CelNavigableMutableExpr::height);
+    // Loop-invariant subexpressions are extracted even if they occur once: block slots are lazily
+    // memoized, so they then run at most once and only if reached, instead of once per iteration.
+    @SuppressWarnings("ValueClassIdentity") // Integer values are only unboxed, never compared.
+    IdentityHashMap<CelMutableExpr, Integer> loopDepths = new IdentityHashMap<>();
+    for (CelNavigableMutableExpr node : allNodes) {
+      if (node.getKind().equals(Kind.COMPREHENSION)) {
+        CelNavigableMutableExpr.fromExpr(node.expr().comprehension().loopStep())
+            .allNodes()
+            .forEach(step -> loopDepths.merge(step.expr(), 1, Integer::sum));
+      }
+    }
+    // Only hoisting needs these, and computing them walks every node's subtree.
+    Set<CelMutableExpr> eligibleExprs = Sets.newIdentityHashSet();
+    Multiset<Integer> vnCounts = HashMultiset.create();
+    if (!loopDepths.isEmpty()) {
+      for (CelNavigableMutableExpr node : allNodes) {
+        if (canEliminate(node, ineligibleExprs)) {
+          eligibleExprs.add(node.expr());
+          vnCounts.add(valueNumberer.vnOf(node.expr()));
+        }
+      }
+    }
     Set<Integer> seenVns = new HashSet<>();
     for (CelNavigableMutableExpr node : allNodes) {
-      ImmutableList<CelNavigableMutableExpr> sameHeightNodes = nodesByHeight.get(node.height());
-      if (sameHeightNodes.size() < 2 || !canEliminate(node, ineligibleExprs)) {
+      boolean inLoopStep = loopDepths.containsKey(node.expr());
+      if ((!inLoopStep && nodesByHeight.get(node.height()).size() < 2)
+          || !canEliminate(node, ineligibleExprs)) {
         continue;
       }
-      int vn = valueNumberer.vnOf(node.expr());
-      if (!seenVns.add(vn)) {
-        return sameHeightNodes.stream()
+      CelNavigableMutableExpr target =
+          inLoopStep ? hoistTarget(node, loopDepths, vnCounts, valueNumberer, eligibleExprs) : node;
+      int vn = valueNumberer.vnOf(target.expr());
+      if (!seenVns.add(vn) || inLoopStep) {
+        return nodesByHeight.get(target.height()).stream()
             .filter(candidate -> canEliminate(candidate, ineligibleExprs))
             .map(CelNavigableMutableExpr::expr)
             .filter(expr -> valueNumberer.vnOf(expr) == vn)
@@ -433,6 +460,41 @@ private ImmutableList<CelMutableExpr> getCseCandidates(
     return ImmutableList.of();
   }
 
+  // Slots are never revisited, so before hoisting `node` this first hoists any eligible descendant
+  // that also occurs outside of it, or that is invariant to a loop nested inside it.
+  private static CelNavigableMutableExpr hoistTarget(
+      CelNavigableMutableExpr node,
+      IdentityHashMap<CelMutableExpr, Integer> loopDepths,
+      Multiset<Integer> vnCounts,
+      ValueNumberer valueNumberer,
+      Set<CelMutableExpr> eligibleExprs) {
+    // A descendant occurring more often than `node` is duplicated within it or shared outside it.
+    int copies = vnCounts.count(valueNumberer.vnOf(node.expr()));
+    int depth = loopDepths.get(node.expr());
+    return node.descendants(TraversalOrder.PRE_ORDER)
+        .filter(
+            d ->
+                eligibleExprs.contains(d.expr())
+                    && (vnCounts.count(valueNumberer.vnOf(d.expr())) > copies
+                        || loopDepths.get(d.expr()) > depth))
+        .findFirst()
+        .map(d -> hoistTarget(d, loopDepths, vnCounts, valueNumberer, eligibleExprs))
+        .orElse(node);
+  }
+
+  // The planner accumulates `@ac + [elem]` steps (map, filter) in place, so `[elem]` is not
+  // extracted.
+  private static boolean isAccumulation(CelMutableExpr expr) {
+    if (!expr.getKind().equals(Kind.CALL)
+        || !expr.call().function().equals(Operator.ADD.getFunction())) {
+      return false;
+    }
+    CelMutableExpr accu = expr.call().args().get(0);
+    return accu.getKind().equals(Kind.IDENT)
+        && accu.ident().name().startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX)
+        && expr.call().args().get(1).getKind().equals(Kind.LIST);
+  }
+
   private boolean canEliminate(
       CelNavigableMutableExpr navigableExpr, Set<CelMutableExpr> ineligibleExprs) {
     return !navigableExpr.getKind().equals(Kind.CONSTANT)
@@ -488,6 +550,12 @@ private static Set<CelMutableExpr> getIneligibleExprsFromComprehensionBranches(
                             .allNodes())
                     .map(CelNavigableMutableExpr::expr)
                     .forEach(ineligibleExprs::add));
+    navAst
+        .getRoot()
+        .allNodes()
+        .map(CelNavigableMutableExpr::expr)
+        .filter(SubexpressionOptimizer::isAccumulation)
+        .forEach(expr -> ineligibleExprs.add(expr.call().args().get(1)));
 
     return ineligibleExprs;
   }

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

@@ -17,6 +17,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static dev.cel.common.CelOverloadDecl.newGlobalOverload;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.testing.junit.testparameterinjector.TestParameter;
 import com.google.testing.junit.testparameterinjector.TestParameterInjector;
@@ -95,6 +96,8 @@ private static Cel setupCelEnv(CelBuilder celBuilder) {
       TestAllTypes.newBuilder()
           .setSingleInt64(3L)
           .setSingleInt32(5)
+          .addAllRepeatedInt64(ImmutableList.of(1L, 2L, 3L))
+          .addRepeatedInt32(2)
           .setOneofType(
               NestedTestAllTypes.newBuilder()
                   .setPayload(
@@ -412,6 +415,26 @@ private enum CseTestCase {
         "(msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 !="
             + " msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 !="
             + " msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)"),
+    HOIST_NESTED_COMPREHENSION("msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == 2))"),
+    HOIST_INVARIANT_ARITHMETIC(
+        "msg.repeated_int64.exists(i, i == msg.single_int64 * 2 + msg.single_int32)"),
+    HOIST_ZERO_TRIP_ERROR_NOT_EVALUATED(
+        "msg.repeated_nested_message.exists(m, m.bb == 1 / (msg.single_int32 - 5))"),
+    HOIST_MAP_ACCUMULATION_PRESERVED("msg.repeated_int64.map(i, msg.single_int64 * 2)"),
+    HOIST_LIST_LITERAL("msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])"),
+    HOIST_LIST_APPENDED_TO_ITER_VAR("[[1], [2]].map(l, l + [msg.single_int64])"),
+    HOIST_INVARIANT_INSIDE_DUPLICATE(
+        "msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32]) ||"
+            + " msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])"),
+    HOIST_INVARIANT_OF_NESTED_COMPREHENSION(
+        "msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == msg.single_int64 * 2))"),
+    HOIST_SHARED_TERMS_ONCE(
+        "msg.repeated_int64.exists(i, i == msg.single_int64 + msg.single_int32 +"
+            + " msg.single_sint64) && [msg.single_int64, msg.single_int32, msg.single_sint64]"
+            + ".size() > 0"),
+    HOIST_SHARED_WITH_ACCUMULATED_LIST(
+        "msg.repeated_int64.exists(i, i in [msg.single_int64]) &&"
+            + " msg.repeated_int64.map(i, msg.single_int64).size() > 0"),
     ;
 
     private final String source;

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

@@ -347,6 +347,27 @@ public void iterationLimitReached_throws() throws Exception {
     assertThat(e).hasMessageThat().isEqualTo("Optimization failure: Max iteration count reached.");
   }
 
+  @Test
+  @TestParameters(
+      "{source: 'msg.repeated_int64.exists(i, i == (msg.single_int64 + msg.single_int32) *"
+          + " (msg.single_int32 + msg.single_int64))', expected: 'cel.@block([msg.single_int64 +"
+          + " msg.single_int32, @index0 * @index0], msg.repeated_int64.exists(@it:0:0, @it:0:0 =="
+          + " @index1))'}")
+  @TestParameters(
+      "{source: 'msg.repeated_int64.map(i, cel.bind(a, i, a + msg.single_int32 * 2))', expected:"
+          + " 'cel.@block([msg.single_int32 * 2], msg.repeated_int64.map(@it:1:0, cel.bind(a,"
+          + " @it:1:0, @ac:0:0 + @index0)))'}")
+  public void cse_hoistLoopInvariants(String source, String expected) throws Exception {
+    CelAbstractSyntaxTree ast = cel.compile(source).getAst();
+    CelOptimizer celOptimizer =
+        newCseOptimizer(
+            SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build());
+
+    CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
+
+    assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(expected);
+  }
+
   @Test
   public void cse_comprehensionReferencingOuterBindVar_notExtractedOutOfBindScope()
       throws Exception {
@@ -757,12 +778,12 @@ public void cse_nestedMacro_noOp_assertAstIdCorrectness() throws Exception {
             .addAstOptimizers(SubexpressionOptimizer.getInstance())
             .build();
     CelAbstractSyntaxTree ast =
-        cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))").getAst();
+        cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))").getAst();
 
     CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
 
     assertThat(CEL_UNPARSER.unparse(optimizedAst))
-        .isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))");
+        .isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))");
     assertThat(optimizedAst).isSameInstanceAs(ast);
   }
 
@@ -797,8 +818,8 @@ public void cse_nestedMacro_withOptimization_assertAstIdCorrectness() throws Exc
 
     assertThat(CEL_UNPARSER.unparse(optimizedAst))
         .isEqualTo(
-            "cel.@block([[{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, has(x.a))], @index0 =="
-                + " @index0)");
+            "cel.@block([has(x.a), [{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, @index0)], @index1"
+                + " == @index1)");
   }
 
   @Test

optimizer/src/test/resources/constfold_before_subexpression_unparsed.baseline

@@ -326,4 +326,64 @@ Test case: VALUE_NUMBERING_COMMUTATIVE_MULTIPLY_EQUALITY_LOGICAL_OR
 Source: (msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 != msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 != msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)
 =====>
 Result: true
-Unparsed: cel.@block([msg.oneof_type.payload, @index0.single_int32 * msg.single_int64 == 24 || msg.single_int32 != @index0.single_int64], @index1 && @index1)
+Unparsed: cel.@block([msg.oneof_type.payload, @index0.single_int32 * msg.single_int64 == 24 || msg.single_int32 != @index0.single_int64], @index1 && @index1)
+
+Test case: HOIST_NESTED_COMPREHENSION
+Source: msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == 2))
+=====>
+Result: true
+Unparsed: cel.@block([msg.repeated_int32.exists(@it:0:0, @it:0:0 == 2)], msg.repeated_int64.all(@it:1:0, @index0))
+
+Test case: HOIST_INVARIANT_ARITHMETIC
+Source: msg.repeated_int64.exists(i, i == msg.single_int64 * 2 + msg.single_int32)
+=====>
+Result: false
+Unparsed: cel.@block([msg.single_int64 * 2 + msg.single_int32], msg.repeated_int64.exists(@it:0:0, @it:0:0 == @index0))
+
+Test case: HOIST_ZERO_TRIP_ERROR_NOT_EVALUATED
+Source: msg.repeated_nested_message.exists(m, m.bb == 1 / (msg.single_int32 - 5))
+=====>
+Result: false
+Unparsed: cel.@block([1 / (msg.single_int32 - 5)], msg.repeated_nested_message.exists(@it:0:0, @it:0:0.bb == @index0))
+
+Test case: HOIST_MAP_ACCUMULATION_PRESERVED
+Source: msg.repeated_int64.map(i, msg.single_int64 * 2)
+=====>
+Result: [6, 6, 6]
+Unparsed: cel.@block([msg.single_int64 * 2], msg.repeated_int64.map(@it:0:0, @index0))
+
+Test case: HOIST_LIST_LITERAL
+Source: msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])
+=====>
+Result: true
+Unparsed: cel.@block([[msg.single_int64, msg.single_int32]], msg.repeated_int64.exists(@it:0:0, @it:0:0 in @index0))
+
+Test case: HOIST_LIST_APPENDED_TO_ITER_VAR
+Source: [[1], [2]].map(l, l + [msg.single_int64])
+=====>
+Result: [[1, 3], [2, 3]]
+Unparsed: cel.@block([[msg.single_int64]], [[1], [2]].map(@it:0:0, @it:0:0 + @index0))
+
+Test case: HOIST_INVARIANT_INSIDE_DUPLICATE
+Source: msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32]) || msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])
+=====>
+Result: true
+Unparsed: cel.@block([[msg.single_int64, msg.single_int32], msg.repeated_int64.exists(@it:0:0, @it:0:0 in @index0)], @index1 || @index1)
+
+Test case: HOIST_INVARIANT_OF_NESTED_COMPREHENSION
+Source: msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == msg.single_int64 * 2))
+=====>
+Result: false
+Unparsed: cel.@block([msg.single_int64 * 2, msg.repeated_int32.exists(@it:0:0, @it:0:0 == @index0)], msg.repeated_int64.all(@it:1:0, @index1))
+
+Test case: HOIST_SHARED_TERMS_ONCE
+Source: msg.repeated_int64.exists(i, i == msg.single_int64 + msg.single_int32 + msg.single_sint64) && [msg.single_int64, msg.single_int32, msg.single_sint64].size() > 0
+=====>
+Result: false
+Unparsed: cel.@block([msg.single_int64, msg.single_int32, msg.single_sint64, @index0 + @index1 + @index2], msg.repeated_int64.exists(@it:0:0, @it:0:0 == @index3) && [@index0, @index1, @index2].size() > 0)
+
+Test case: HOIST_SHARED_WITH_ACCUMULATED_LIST
+Source: msg.repeated_int64.exists(i, i in [msg.single_int64]) && msg.repeated_int64.map(i, msg.single_int64).size() > 0
+=====>
+Result: true
+Unparsed: cel.@block([msg.single_int64, [@index0], msg.repeated_int64], @index2.exists(@it:0:0, @it:0:0 in @index1) && @index2.map(@it:0:1, @index0).size() > 0)