Skip to content

Commit 7510662

Browse files
l46kokcopybara-github
authored andcommitted
Add value numbering to SubexpressionOptimizer for commutative operators and relational converses
PiperOrigin-RevId: 1003358263
1 parent 3c5d875 commit 7510662

7 files changed

Lines changed: 515 additions & 22 deletions

File tree

‎optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,11 +60,13 @@ java_library(
6060
":default_optimizer_constants",
6161
"//:auto_value",
6262
"//bundle:cel",
63+
"//checker:standard_decl",
6364
"//common:cel_ast",
6465
"//common:cel_source",
6566
"//common:compiler_common",
6667
"//common:mutable_ast",
6768
"//common:mutable_source",
69+
"//common:operator",
6870
"//common/ast",
6971
"//common/ast:cel_block",
7072
"//common/ast:mutable_expr",

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

Lines changed: 164 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,18 @@
2828
import com.google.common.base.Verify;
2929
import com.google.common.collect.ImmutableList;
3030
import com.google.common.collect.ImmutableListMultimap;
31+
import com.google.common.collect.ImmutableMap;
3132
import com.google.common.collect.ImmutableSet;
3233
import com.google.common.collect.Multimaps;
3334
import com.google.common.collect.Sets;
3435
import com.google.common.collect.Streams;
3536
import com.google.errorprone.annotations.CanIgnoreReturnValue;
3637
import dev.cel.bundle.Cel;
3738
import dev.cel.bundle.CelBuilder;
39+
import dev.cel.checker.CelStandardDeclarations.StandardFunction.Overload.Arithmetic;
40+
import dev.cel.checker.CelStandardDeclarations.StandardFunction.Overload.BooleanOperator;
41+
import dev.cel.checker.CelStandardDeclarations.StandardFunction.Overload.Comparison;
42+
import dev.cel.checker.CelStandardDeclarations.StandardFunction.Overload.Relation;
3843
import dev.cel.common.CelAbstractSyntaxTree;
3944
import dev.cel.common.CelMutableAst;
4045
import dev.cel.common.CelMutableSource;
@@ -44,15 +49,23 @@
4449
import dev.cel.common.CelSource.Extension.Version;
4550
import dev.cel.common.CelValidationException;
4651
import dev.cel.common.CelVarDecl;
52+
import dev.cel.common.Operator;
4753
import dev.cel.common.ast.CelBlock;
4854
import dev.cel.common.ast.CelExpr;
4955
import dev.cel.common.ast.CelExpr.CelCall;
5056
import dev.cel.common.ast.CelExpr.CelComprehension;
5157
import dev.cel.common.ast.CelExpr.CelList;
5258
import dev.cel.common.ast.CelExpr.ExprKind.Kind;
5359
import dev.cel.common.ast.CelMutableExpr;
60+
import dev.cel.common.ast.CelMutableExpr.CelMutableCall;
61+
import dev.cel.common.ast.CelMutableExpr.CelMutableComprehension;
5462
import dev.cel.common.ast.CelMutableExpr.CelMutableIdent;
63+
import dev.cel.common.ast.CelMutableExpr.CelMutableList;
64+
import dev.cel.common.ast.CelMutableExpr.CelMutableMap;
65+
import dev.cel.common.ast.CelMutableExpr.CelMutableSelect;
66+
import dev.cel.common.ast.CelMutableExpr.CelMutableStruct;
5567
import dev.cel.common.ast.CelMutableExprConverter;
68+
import dev.cel.common.ast.CelReference;
5669
import dev.cel.common.navigation.CelNavigableExpr;
5770
import dev.cel.common.navigation.CelNavigableExprUtil;
5871
import dev.cel.common.navigation.CelNavigableMutableAst;
@@ -65,8 +78,10 @@
6578
import dev.cel.optimizer.CelAstOptimizer;
6679
import java.util.ArrayList;
6780
import java.util.Arrays;
81+
import java.util.EnumSet;
6882
import java.util.HashMap;
6983
import java.util.HashSet;
84+
import java.util.IdentityHashMap;
7085
import java.util.List;
7186
import java.util.Map;
7287
import java.util.Set;
@@ -173,7 +188,9 @@ private OptimizationResult optimizeUsingCelBlock(CelAbstractSyntaxTree ast, Cel
173188

174189
for (iterCount = 0; iterCount < cseOptions.iterationLimit(); iterCount++) {
175190
CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(astToModify);
176-
ImmutableList<CelMutableExpr> cseCandidates = getCseCandidates(navAst);
191+
ValueNumberer valueNumberer =
192+
new ValueNumberer(navAst.getRoot().expr(), ast.getReferenceMap());
193+
ImmutableList<CelMutableExpr> cseCandidates = getCseCandidates(navAst, valueNumberer);
177194
if (cseCandidates.isEmpty()) {
178195
break;
179196
}
@@ -390,25 +407,25 @@ private static void replaceWithBlockIdent(
390407
expr.setIdent(CelMutableIdent.create(blockIdentifier));
391408
}
392409

393-
private ImmutableList<CelMutableExpr> getCseCandidates(CelNavigableMutableAst navAst) {
410+
private ImmutableList<CelMutableExpr> getCseCandidates(
411+
CelNavigableMutableAst navAst, ValueNumberer valueNumberer) {
394412
Set<CelMutableExpr> ineligibleExprs = getIneligibleExprsFromComprehensionBranches(navAst);
395413
ImmutableList<CelNavigableMutableExpr> allNodes =
396414
navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).collect(toImmutableList());
397415
ImmutableListMultimap<Integer, CelNavigableMutableExpr> nodesByHeight =
398416
Multimaps.index(allNodes, CelNavigableMutableExpr::height);
399-
HashSet<CelMutableExpr> semanticallyEqualNodes = new HashSet<>();
417+
Set<Integer> seenVns = new HashSet<>();
400418
for (CelNavigableMutableExpr node : allNodes) {
401419
ImmutableList<CelNavigableMutableExpr> sameHeightNodes = nodesByHeight.get(node.height());
402420
if (sameHeightNodes.size() < 2 || !canEliminate(node, ineligibleExprs)) {
403421
continue;
404422
}
405-
// Normalize the expr to test semantic equivalence.
406-
CelMutableExpr normalizedExpr = normalizeForEquality(node.expr());
407-
if (!semanticallyEqualNodes.add(normalizedExpr)) {
423+
int vn = valueNumberer.vnOf(node.expr());
424+
if (!seenVns.add(vn)) {
408425
return sameHeightNodes.stream()
409426
.filter(candidate -> canEliminate(candidate, ineligibleExprs))
410427
.map(CelNavigableMutableExpr::expr)
411-
.filter(expr -> normalizeForEquality(expr).equals(normalizedExpr))
428+
.filter(expr -> valueNumberer.vnOf(expr) == vn)
412429
.collect(toImmutableList());
413430
}
414431
}
@@ -488,17 +505,147 @@ private boolean containsEliminableFunctionOnly(CelNavigableMutableExpr navigable
488505
});
489506
}
490507

491-
/**
492-
* Converts the {@link CelMutableExpr} to make it suitable for performing a semantically equals
493-
* check.
494-
*
495-
* <p>Specifically, this will deep copy the mutable expr then set all expr IDs in the expression
496-
* tree to 0.
497-
*/
498-
private CelMutableExpr normalizeForEquality(CelMutableExpr mutableExpr) {
499-
CelMutableExpr copiedExpr = CelMutableExpr.newInstance(mutableExpr);
508+
// Gives equal value numbers to subexpressions that always compute the same value: identical
509+
// subtrees, commuted calls (`a + b`, `b + a`) and converse comparisons (`a > b`, `b < a`). Calls
510+
// are only reordered if the checker resolved them to standard overloads, as custom overloads may
511+
// not be commutative. If both operands fail, the merged call may report the other operand's
512+
// error; CEL leaves evaluation order, and thus which error is reported, unspecified.
513+
private static final class ValueNumberer {
514+
private static final ImmutableMap<String, String> SWAPPED_FUNCTIONS =
515+
ImmutableMap.<String, String>builder()
516+
.put(Operator.ADD.getFunction(), Operator.ADD.getFunction())
517+
.put(Operator.MULTIPLY.getFunction(), Operator.MULTIPLY.getFunction())
518+
.put(Operator.EQUALS.getFunction(), Operator.EQUALS.getFunction())
519+
.put(Operator.NOT_EQUALS.getFunction(), Operator.NOT_EQUALS.getFunction())
520+
.put(Operator.LOGICAL_AND.getFunction(), Operator.LOGICAL_AND.getFunction())
521+
.put(Operator.LOGICAL_OR.getFunction(), Operator.LOGICAL_OR.getFunction())
522+
.put(Operator.GREATER.getFunction(), Operator.LESS.getFunction())
523+
.put(Operator.GREATER_EQUALS.getFunction(), Operator.LESS_EQUALS.getFunction())
524+
.buildOrThrow();
525+
private static final ImmutableSet<String> SWAPPABLE_OVERLOAD_IDS =
526+
Stream.concat(
527+
Stream.of(
528+
Arithmetic.ADD_INT64,
529+
Arithmetic.ADD_UINT64,
530+
Arithmetic.ADD_DOUBLE,
531+
Arithmetic.ADD_DURATION_DURATION,
532+
Arithmetic.ADD_TIMESTAMP_DURATION,
533+
Arithmetic.ADD_DURATION_TIMESTAMP,
534+
Arithmetic.MULTIPLY_INT64,
535+
Arithmetic.MULTIPLY_UINT64,
536+
Arithmetic.MULTIPLY_DOUBLE,
537+
Relation.EQUALS,
538+
Relation.NOT_EQUALS,
539+
BooleanOperator.LOGICAL_AND,
540+
BooleanOperator.LOGICAL_OR),
541+
EnumSet.allOf(Comparison.class).stream())
542+
.map(overload -> overload.celOverloadDecl().overloadId())
543+
.collect(toImmutableSet());
544+
545+
private final ImmutableMap<Long, CelReference> references;
546+
private final Map<ImmutableList<?>, Integer> keyToVn;
547+
private final IdentityHashMap<CelMutableExpr, Integer> exprToVn;
548+
549+
private int vnOf(CelMutableExpr expr) {
550+
return Verify.verifyNotNull(exprToVn.get(expr));
551+
}
552+
553+
private int computeVn(CelMutableExpr expr) {
554+
ImmutableList<?> key = buildKey(expr);
555+
int vn = keyToVn.computeIfAbsent(key, k -> keyToVn.size());
556+
exprToVn.put(expr, vn);
557+
return vn;
558+
}
500559

501-
return astMutator.clearExprIds(copiedExpr);
560+
private ImmutableList<Integer> computeVns(List<CelMutableExpr> exprs) {
561+
ImmutableList.Builder<Integer> vns = ImmutableList.builderWithExpectedSize(exprs.size());
562+
for (CelMutableExpr expr : exprs) {
563+
vns.add(computeVn(expr));
564+
}
565+
return vns.build();
566+
}
567+
568+
private ImmutableList<?> buildKey(CelMutableExpr expr) {
569+
switch (expr.getKind()) {
570+
case CONSTANT:
571+
return ImmutableList.of(Kind.CONSTANT, expr.constant());
572+
case IDENT:
573+
return ImmutableList.of(Kind.IDENT, expr.ident().name());
574+
case SELECT:
575+
CelMutableSelect select = expr.select();
576+
return ImmutableList.of(
577+
Kind.SELECT, computeVn(select.operand()), select.field(), select.testOnly());
578+
case CALL:
579+
CelMutableCall call = expr.call();
580+
String function = call.function();
581+
// VNs are non-negative, so -1 marks a global call.
582+
int targetVn = call.target().map(this::computeVn).orElse(-1);
583+
ImmutableList<Integer> argVns = computeVns(call.args());
584+
// Keyed even when no swap applies (e.g. _<_), so a call that also resolves to custom
585+
// overloads (x < 5 on dyn) never merges with a canonicalized standard-only one (5 > x).
586+
boolean swappable = hasOnlySwappableOverloads(expr.id());
587+
String swapped = SWAPPED_FUNCTIONS.get(function);
588+
if (swappable
589+
&& swapped != null
590+
&& (!swapped.equals(function) || argVns.get(0) > argVns.get(1))) {
591+
function = swapped;
592+
argVns = argVns.reverse();
593+
}
594+
return ImmutableList.of(Kind.CALL, targetVn, function, argVns, swappable);
595+
case LIST:
596+
CelMutableList list = expr.list();
597+
return ImmutableList.of(
598+
Kind.LIST, computeVns(list.elements()), ImmutableList.copyOf(list.optionalIndices()));
599+
case MAP:
600+
CelMutableMap map = expr.map();
601+
ImmutableList.Builder<Object> entryVns =
602+
ImmutableList.builderWithExpectedSize(map.entries().size());
603+
for (CelMutableMap.Entry entry : map.entries()) {
604+
entryVns.add(
605+
ImmutableList.of(
606+
computeVn(entry.key()), computeVn(entry.value()), entry.optionalEntry()));
607+
}
608+
return ImmutableList.of(Kind.MAP, entryVns.build());
609+
case STRUCT:
610+
CelMutableStruct struct = expr.struct();
611+
ImmutableList.Builder<Object> fieldVns =
612+
ImmutableList.builderWithExpectedSize(struct.entries().size());
613+
for (CelMutableStruct.Entry entry : struct.entries()) {
614+
fieldVns.add(
615+
ImmutableList.of(
616+
entry.fieldKey(), computeVn(entry.value()), entry.optionalEntry()));
617+
}
618+
return ImmutableList.of(Kind.STRUCT, struct.messageName(), fieldVns.build());
619+
case COMPREHENSION:
620+
CelMutableComprehension comp = expr.comprehension();
621+
return ImmutableList.of(
622+
Kind.COMPREHENSION,
623+
comp.iterVar(),
624+
comp.iterVar2(),
625+
comp.accuVar(),
626+
computeVn(comp.iterRange()),
627+
computeVn(comp.accuInit()),
628+
computeVn(comp.loopCondition()),
629+
computeVn(comp.loopStep()),
630+
computeVn(comp.result()));
631+
default:
632+
throw new IllegalArgumentException("Unexpected ExprKind: " + expr.getKind());
633+
}
634+
}
635+
636+
private boolean hasOnlySwappableOverloads(long id) {
637+
ImmutableList<String> overloadIds = references.get(id).overloadIds();
638+
// Fail closed: containsAll is vacuously true for an empty list.
639+
return !overloadIds.isEmpty() && SWAPPABLE_OVERLOAD_IDS.containsAll(overloadIds);
640+
}
641+
642+
@SuppressWarnings("ValueClassIdentity") // Integer values are only unboxed, never compared.
643+
private ValueNumberer(CelMutableExpr root, ImmutableMap<Long, CelReference> references) {
644+
this.references = references;
645+
this.keyToVn = new HashMap<>();
646+
this.exprToVn = new IdentityHashMap<>();
647+
computeVn(root);
648+
}
502649
}
503650

504651
/** Options to configure how Common Subexpression Elimination behave. */

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -403,7 +403,17 @@ private enum CseTestCase {
403403
"pure_custom_func(msg.oneof_type.payload.single_int64) +"
404404
+ " pure_custom_func(msg.oneof_type.payload.single_int32) +"
405405
+ " pure_custom_func(msg.oneof_type.payload.single_int64) +"
406-
+ " pure_custom_func(msg.single_int64)");
406+
+ " pure_custom_func(msg.single_int64)"),
407+
VALUE_NUMBERING_COMMUTATIVE_ADD_CONVERSE_RELATION_LOGICAL_AND(
408+
"(msg.single_int64 + msg.single_int32 > 5 && msg.oneof_type.payload.single_int64 <"
409+
+ " msg.single_int64 + msg.single_int32) == (msg.single_int32 + msg.single_int64 >"
410+
+ " msg.oneof_type.payload.single_int64 && 5 < msg.single_int32 + msg.single_int64)"),
411+
VALUE_NUMBERING_COMMUTATIVE_MULTIPLY_EQUALITY_LOGICAL_OR(
412+
"(msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 !="
413+
+ " msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 !="
414+
+ " msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)"),
415+
;
416+
407417
private final String source;
408418

409419
CseTestCase(String source) {

0 commit comments

Comments
 (0)