Skip to content

Commit 2e1f5a2

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Extract shared parts of duplicates first in SubexpressionOptimizer
PiperOrigin-RevId: 1005861790
1 parent b7096de commit 2e1f5a2

8 files changed

Lines changed: 2051 additions & 720 deletions

File tree

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

Lines changed: 108 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,11 @@
2626
import com.google.common.annotations.VisibleForTesting;
2727
import com.google.common.base.Strings;
2828
import com.google.common.base.Verify;
29+
import com.google.common.collect.HashMultiset;
2930
import com.google.common.collect.ImmutableList;
30-
import com.google.common.collect.ImmutableListMultimap;
3131
import com.google.common.collect.ImmutableMap;
3232
import com.google.common.collect.ImmutableSet;
33-
import com.google.common.collect.Multimaps;
33+
import com.google.common.collect.Multiset;
3434
import com.google.common.collect.Sets;
3535
import com.google.common.collect.Streams;
3636
import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -85,7 +85,6 @@
8585
import java.util.List;
8686
import java.util.Map;
8787
import java.util.Set;
88-
import java.util.function.Predicate;
8988
import java.util.stream.Stream;
9089

9190
/**
@@ -412,55 +411,118 @@ private ImmutableList<CelMutableExpr> getCseCandidates(
412411
Set<CelMutableExpr> ineligibleExprs = getIneligibleExprsFromComprehensionBranches(navAst);
413412
ImmutableList<CelNavigableMutableExpr> allNodes =
414413
navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).collect(toImmutableList());
415-
ImmutableListMultimap<Integer, CelNavigableMutableExpr> nodesByHeight =
416-
Multimaps.index(allNodes, CelNavigableMutableExpr::height);
417-
Set<Integer> seenVns = new HashSet<>();
414+
// Loop-invariant subexpressions are extracted even if they occur once: block slots are lazily
415+
// memoized, so they then run at most once and only if reached, instead of once per iteration.
416+
@SuppressWarnings("ValueClassIdentity") // Integer values are only unboxed, never compared.
417+
IdentityHashMap<CelMutableExpr, Integer> loopDepths = new IdentityHashMap<>();
418418
for (CelNavigableMutableExpr node : allNodes) {
419-
ImmutableList<CelNavigableMutableExpr> sameHeightNodes = nodesByHeight.get(node.height());
420-
if (sameHeightNodes.size() < 2 || !canEliminate(node, ineligibleExprs)) {
421-
continue;
419+
if (node.getKind().equals(Kind.COMPREHENSION)) {
420+
CelNavigableMutableExpr.fromExpr(node.expr().comprehension().loopStep())
421+
.allNodes()
422+
.forEach(step -> loopDepths.merge(step.expr(), 1, Integer::sum));
422423
}
423-
int vn = valueNumberer.vnOf(node.expr());
424-
if (!seenVns.add(vn)) {
425-
return sameHeightNodes.stream()
426-
.filter(candidate -> canEliminate(candidate, ineligibleExprs))
424+
}
425+
Set<CelMutableExpr> eligibleExprs = getEligibleExprs(allNodes, ineligibleExprs);
426+
Multiset<Integer> vnCounts = HashMultiset.create();
427+
eligibleExprs.forEach(expr -> vnCounts.add(valueNumberer.vnOf(expr)));
428+
Set<Integer> seenVns = new HashSet<>();
429+
for (CelNavigableMutableExpr node : allNodes) {
430+
CelMutableExpr expr = node.expr();
431+
if (eligibleExprs.contains(expr)
432+
&& (loopDepths.containsKey(expr) || !seenVns.add(valueNumberer.vnOf(expr)))) {
433+
int vn =
434+
valueNumberer.vnOf(
435+
extractionTarget(node, loopDepths, vnCounts, valueNumberer, eligibleExprs).expr());
436+
return allNodes.stream()
427437
.map(CelNavigableMutableExpr::expr)
428-
.filter(expr -> valueNumberer.vnOf(expr) == vn)
438+
.filter(e -> eligibleExprs.contains(e) && valueNumberer.vnOf(e) == vn)
429439
.collect(toImmutableList());
430440
}
431441
}
432442

433443
return ImmutableList.of();
434444
}
435445

436-
private boolean canEliminate(
437-
CelNavigableMutableExpr navigableExpr, Set<CelMutableExpr> ineligibleExprs) {
438-
return !navigableExpr.getKind().equals(Kind.CONSTANT)
439-
&& !navigableExpr.getKind().equals(Kind.IDENT)
440-
// Exclude empty lists (cel.bind sets this for iterRange).
441-
&& !(navigableExpr.getKind().equals(Kind.LIST)
442-
&& navigableExpr.expr().list().elements().isEmpty())
443-
&& !ineligibleExprs.contains(navigableExpr.expr())
444-
&& containsEliminableFunctionOnly(navigableExpr)
445-
&& containsValidComprehensionIdents(navigableExpr);
446+
// Slots are never revisited, so before extracting `node` this first extracts any eligible
447+
// descendant that is duplicated within or also occurs outside of it, or (for a loop-step node)
448+
// that is invariant to a loop nested inside it.
449+
private static CelNavigableMutableExpr extractionTarget(
450+
CelNavigableMutableExpr node,
451+
IdentityHashMap<CelMutableExpr, Integer> loopDepths,
452+
Multiset<Integer> vnCounts,
453+
ValueNumberer valueNumberer,
454+
Set<CelMutableExpr> eligibleExprs) {
455+
// A descendant occurring more often than `node` is duplicated within it or shared outside it.
456+
int copies = vnCounts.count(valueNumberer.vnOf(node.expr()));
457+
Integer depth = loopDepths.get(node.expr());
458+
return node.descendants(TraversalOrder.PRE_ORDER)
459+
.filter(
460+
d ->
461+
eligibleExprs.contains(d.expr())
462+
&& (vnCounts.count(valueNumberer.vnOf(d.expr())) > copies
463+
|| (depth != null && loopDepths.get(d.expr()) > depth)))
464+
.findFirst()
465+
.map(d -> extractionTarget(d, loopDepths, vnCounts, valueNumberer, eligibleExprs))
466+
.orElse(node);
446467
}
447468

448-
private static boolean containsValidComprehensionIdents(CelNavigableMutableExpr navExpr) {
449-
// For nested comprehensions of form [1].exists(x, [2].exists(y, x == y)), the inner
450-
// comprehension [2].exists(y, x == y) should not be extracted out into a block index, as it
451-
// causes issues with scoping. Non-comprehension subexpressions cannot reference any mangled
452-
// comprehension variables.
453-
Predicate<String> isInvalidIdent =
454-
navExpr.getKind().equals(Kind.COMPREHENSION)
455-
? CelNavigableExprUtil.getEnclosingComprehensionVariables(navExpr)::contains
456-
: name ->
457-
name.startsWith(MANGLED_COMPREHENSION_ITER_VAR_PREFIX)
458-
|| name.startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX);
459-
return navExpr
460-
.descendants()
461-
.filter(node -> node.getKind().equals(Kind.IDENT))
462-
.map(node -> node.expr().ident().name())
463-
.noneMatch(isInvalidIdent);
469+
// The planner accumulates `@ac + [elem]` steps (map, filter) in place, so `[elem]` is not
470+
// extracted.
471+
private static boolean isAccumulation(CelMutableExpr expr) {
472+
if (!expr.getKind().equals(Kind.CALL)
473+
|| !expr.call().function().equals(Operator.ADD.getFunction())) {
474+
return false;
475+
}
476+
CelMutableExpr accu = expr.call().args().get(0);
477+
return accu.getKind().equals(Kind.IDENT)
478+
&& accu.ident().name().startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX)
479+
&& expr.call().args().get(1).getKind().equals(Kind.LIST);
480+
}
481+
482+
// Computed bottom-up, as checking every node's subtree separately is quadratic. Only
483+
// comprehensions still walk their own subtree.
484+
private Set<CelMutableExpr> getEligibleExprs(
485+
ImmutableList<CelNavigableMutableExpr> preOrderNodes, Set<CelMutableExpr> ineligibleExprs) {
486+
// Subtrees containing an ineliminable call, or any mangled comprehension variable.
487+
Set<CelMutableExpr> impureExprs = Sets.newIdentityHashSet();
488+
Set<CelMutableExpr> mangledExprs = Sets.newIdentityHashSet();
489+
Set<CelMutableExpr> eligibleExprs = Sets.newIdentityHashSet();
490+
// Reverse pre-order visits children before their parents.
491+
for (CelNavigableMutableExpr node : preOrderNodes.reverse()) {
492+
CelMutableExpr expr = node.expr();
493+
Kind kind = node.getKind();
494+
if ((kind.equals(Kind.CALL) && !cseEliminableFunctions.contains(expr.call().function()))
495+
|| node.children().anyMatch(child -> impureExprs.contains(child.expr()))) {
496+
impureExprs.add(expr);
497+
}
498+
if ((kind.equals(Kind.IDENT)
499+
&& (expr.ident().name().startsWith(MANGLED_COMPREHENSION_ITER_VAR_PREFIX)
500+
|| expr.ident().name().startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX)))
501+
|| node.children().anyMatch(child -> mangledExprs.contains(child.expr()))) {
502+
mangledExprs.add(expr);
503+
}
504+
if (!kind.equals(Kind.CONSTANT)
505+
&& !kind.equals(Kind.IDENT)
506+
// Exclude empty lists (cel.bind sets this for iterRange).
507+
&& !(kind.equals(Kind.LIST) && expr.list().elements().isEmpty())
508+
&& !ineligibleExprs.contains(expr)
509+
&& !impureExprs.contains(expr)
510+
&& (kind.equals(Kind.COMPREHENSION)
511+
? !referencesEnclosingComprehensionVars(node)
512+
: !mangledExprs.contains(expr))) {
513+
eligibleExprs.add(expr);
514+
}
515+
}
516+
return eligibleExprs;
517+
}
518+
519+
// Comprehensions may reference their own variables, but not those of enclosing ones (e.g.
520+
// [2].exists(y, x == y) in [1].exists(x, [2].exists(y, x == y))), as that breaks scoping.
521+
private static boolean referencesEnclosingComprehensionVars(CelNavigableMutableExpr node) {
522+
return node.descendants()
523+
.filter(d -> d.getKind().equals(Kind.IDENT))
524+
.map(d -> d.expr().ident().name())
525+
.anyMatch(CelNavigableExprUtil.getEnclosingComprehensionVariables(node)::contains);
464526
}
465527

466528
/**
@@ -488,21 +550,14 @@ private static Set<CelMutableExpr> getIneligibleExprsFromComprehensionBranches(
488550
.allNodes())
489551
.map(CelNavigableMutableExpr::expr)
490552
.forEach(ineligibleExprs::add));
491-
492-
return ineligibleExprs;
493-
}
494-
495-
private boolean containsEliminableFunctionOnly(CelNavigableMutableExpr navigableExpr) {
496-
return navigableExpr
553+
navAst
554+
.getRoot()
497555
.allNodes()
498-
.allMatch(
499-
node -> {
500-
if (node.getKind().equals(Kind.CALL)) {
501-
return cseEliminableFunctions.contains(node.expr().call().function());
502-
}
556+
.map(CelNavigableMutableExpr::expr)
557+
.filter(SubexpressionOptimizer::isAccumulation)
558+
.forEach(expr -> ineligibleExprs.add(expr.call().args().get(1)));
503559

504-
return true;
505-
});
560+
return ineligibleExprs;
506561
}
507562

508563
// Assigns equal value numbers to equivalent subexpressions, including commuted and converse calls

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import static com.google.common.truth.Truth.assertThat;
1818
import static dev.cel.common.CelOverloadDecl.newGlobalOverload;
1919

20+
import com.google.common.collect.ImmutableList;
2021
import com.google.common.collect.ImmutableMap;
2122
import com.google.testing.junit.testparameterinjector.TestParameter;
2223
import com.google.testing.junit.testparameterinjector.TestParameterInjector;
@@ -95,6 +96,8 @@ private static Cel setupCelEnv(CelBuilder celBuilder) {
9596
TestAllTypes.newBuilder()
9697
.setSingleInt64(3L)
9798
.setSingleInt32(5)
99+
.addAllRepeatedInt64(ImmutableList.of(1L, 2L, 3L))
100+
.addRepeatedInt32(2)
98101
.setOneofType(
99102
NestedTestAllTypes.newBuilder()
100103
.setPayload(
@@ -412,6 +415,29 @@ private enum CseTestCase {
412415
"(msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 !="
413416
+ " msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 !="
414417
+ " msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)"),
418+
HOIST_NESTED_COMPREHENSION("msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == 2))"),
419+
HOIST_INVARIANT_ARITHMETIC(
420+
"msg.repeated_int64.exists(i, i == msg.single_int64 * 2 + msg.single_int32)"),
421+
HOIST_ZERO_TRIP_ERROR_NOT_EVALUATED(
422+
"msg.repeated_nested_message.exists(m, m.bb == 1 / (msg.single_int32 - 5))"),
423+
HOIST_MAP_ACCUMULATION_PRESERVED("msg.repeated_int64.map(i, msg.single_int64 * 2)"),
424+
HOIST_LIST_LITERAL("msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])"),
425+
HOIST_LIST_APPENDED_TO_ITER_VAR("[[1], [2]].map(l, l + [msg.single_int64])"),
426+
HOIST_INVARIANT_INSIDE_DUPLICATE(
427+
"msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32]) ||"
428+
+ " msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])"),
429+
HOIST_INVARIANT_OF_NESTED_COMPREHENSION(
430+
"msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == msg.single_int64 * 2))"),
431+
HOIST_SHARED_TERMS_ONCE(
432+
"msg.repeated_int64.exists(i, i == msg.single_int64 + msg.single_int32 +"
433+
+ " msg.single_sint64) && [msg.single_int64, msg.single_int32, msg.single_sint64]"
434+
+ ".size() > 0"),
435+
HOIST_SHARED_WITH_ACCUMULATED_LIST(
436+
"msg.repeated_int64.exists(i, i in [msg.single_int64]) &&"
437+
+ " msg.repeated_int64.map(i, msg.single_int64).size() > 0"),
438+
DUPLICATE_WITH_SHARED_COMMUTED_OPERANDS(
439+
"[msg.single_int64 * msg.single_int32, msg.single_int32 * msg.single_int64,"
440+
+ " msg.single_int64 + msg.single_int32, msg.single_int32 + msg.single_int64]"),
415441
;
416442

417443
private final String source;

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

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -302,7 +302,9 @@ public void cse_applyConstFoldingAfter() throws Exception {
302302
CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);
303303

304304
assertThat(CEL_UNPARSER.unparse(optimizedAst))
305-
.isEqualTo("cel.@block([1, 2], @index0 + @index0 + @index1 + @index1 + x)");
305+
.isEqualTo(
306+
"cel.@block([3, size([@index0]), size([1, @index0])], @index1 + @index1 + @index2 +"
307+
+ " @index2 + x)");
306308
}
307309

308310
@Test
@@ -347,6 +349,27 @@ public void iterationLimitReached_throws() throws Exception {
347349
assertThat(e).hasMessageThat().isEqualTo("Optimization failure: Max iteration count reached.");
348350
}
349351

352+
@Test
353+
@TestParameters(
354+
"{source: 'msg.repeated_int64.exists(i, i == (msg.single_int64 + msg.single_int32) *"
355+
+ " (msg.single_int32 + msg.single_int64))', expected: 'cel.@block([msg.single_int64 +"
356+
+ " msg.single_int32, @index0 * @index0], msg.repeated_int64.exists(@it:0:0, @it:0:0 =="
357+
+ " @index1))'}")
358+
@TestParameters(
359+
"{source: 'msg.repeated_int64.map(i, cel.bind(a, i, a + msg.single_int32 * 2))', expected:"
360+
+ " 'cel.@block([msg.single_int32 * 2], msg.repeated_int64.map(@it:1:0, cel.bind(a,"
361+
+ " @it:1:0, @ac:0:0 + @index0)))'}")
362+
public void cse_hoistLoopInvariants(String source, String expected) throws Exception {
363+
CelAbstractSyntaxTree ast = cel.compile(source).getAst();
364+
CelOptimizer celOptimizer =
365+
newCseOptimizer(
366+
SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build());
367+
368+
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
369+
370+
assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(expected);
371+
}
372+
350373
@Test
351374
public void cse_comprehensionReferencingOuterBindVar_notExtractedOutOfBindScope()
352375
throws Exception {
@@ -757,12 +780,12 @@ public void cse_nestedMacro_noOp_assertAstIdCorrectness() throws Exception {
757780
.addAstOptimizers(SubexpressionOptimizer.getInstance())
758781
.build();
759782
CelAbstractSyntaxTree ast =
760-
cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))").getAst();
783+
cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))").getAst();
761784

762785
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
763786

764787
assertThat(CEL_UNPARSER.unparse(optimizedAst))
765-
.isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))");
788+
.isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))");
766789
assertThat(optimizedAst).isSameInstanceAs(ast);
767790
}
768791

@@ -797,8 +820,8 @@ public void cse_nestedMacro_withOptimization_assertAstIdCorrectness() throws Exc
797820

798821
assertThat(CEL_UNPARSER.unparse(optimizedAst))
799822
.isEqualTo(
800-
"cel.@block([[{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, has(x.a))], @index0 =="
801-
+ " @index0)");
823+
"cel.@block([has(x.a), [{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, @index0)], @index1"
824+
+ " == @index1)");
802825
}
803826

804827
@Test
@@ -869,7 +892,7 @@ private enum ValueNumberingTestCase {
869892
"cel.@block([msg.single_int64 <= msg.single_int32], @index0 && @index0)"),
870893
NOT_EQUALS_WITHOUT_SHARED_EQUALS_PRESERVED(
871894
"[msg.single_int64 != 1, msg.single_int64 != 1, msg.single_int64 == 1]",
872-
"cel.@block([msg.single_int64 != 1], [@index0, @index0, msg.single_int64 == 1])"),
895+
"cel.@block([msg.single_int64, @index0 != 1], [@index1, @index1, @index0 == 1])"),
873896
OPTIONAL_LIST_AND_MAP_DISTINCT_VALUE_NUMBERS(
874897
"([?optional.of(msg.single_int64)].size() == [optional.of(msg.single_int64)].size()) &&"
875898
+ " ({?\"k\": optional.of(msg.single_int64)}.size() == {\"k\":"

0 commit comments

Comments
 (0)