Skip to content

Commit dbecb80

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Add value numbering to SubexpressionOptimizer
PiperOrigin-RevId: 1003358263
1 parent b6a64ff commit dbecb80

7 files changed

Lines changed: 514 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: 163 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,146 @@ 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+
// Bottom-up hash-consing value numberer. Binary calls f(a, b) and SWAPPED_FUNCTIONS[f](b, a)
509+
// share a value number when the checker resolved them only to standard overloads, as custom
510+
// overloads need not be commutative or converse. When both operands error, the surfaced error may
511+
// differ, which CEL leaves unspecified.
512+
private static final class ValueNumberer {
513+
private static final ImmutableMap<String, String> SWAPPED_FUNCTIONS =
514+
ImmutableMap.<String, String>builder()
515+
.put(Operator.ADD.getFunction(), Operator.ADD.getFunction())
516+
.put(Operator.MULTIPLY.getFunction(), Operator.MULTIPLY.getFunction())
517+
.put(Operator.EQUALS.getFunction(), Operator.EQUALS.getFunction())
518+
.put(Operator.NOT_EQUALS.getFunction(), Operator.NOT_EQUALS.getFunction())
519+
.put(Operator.LOGICAL_AND.getFunction(), Operator.LOGICAL_AND.getFunction())
520+
.put(Operator.LOGICAL_OR.getFunction(), Operator.LOGICAL_OR.getFunction())
521+
.put(Operator.GREATER.getFunction(), Operator.LESS.getFunction())
522+
.put(Operator.GREATER_EQUALS.getFunction(), Operator.LESS_EQUALS.getFunction())
523+
.buildOrThrow();
524+
private static final ImmutableSet<String> SWAPPABLE_OVERLOAD_IDS =
525+
Stream.concat(
526+
Stream.of(
527+
Arithmetic.ADD_INT64,
528+
Arithmetic.ADD_UINT64,
529+
Arithmetic.ADD_DOUBLE,
530+
Arithmetic.ADD_DURATION_DURATION,
531+
Arithmetic.ADD_TIMESTAMP_DURATION,
532+
Arithmetic.ADD_DURATION_TIMESTAMP,
533+
Arithmetic.MULTIPLY_INT64,
534+
Arithmetic.MULTIPLY_UINT64,
535+
Arithmetic.MULTIPLY_DOUBLE,
536+
Relation.EQUALS,
537+
Relation.NOT_EQUALS,
538+
BooleanOperator.LOGICAL_AND,
539+
BooleanOperator.LOGICAL_OR),
540+
EnumSet.allOf(Comparison.class).stream())
541+
.map(overload -> overload.celOverloadDecl().overloadId())
542+
.collect(toImmutableSet());
543+
544+
private final ImmutableMap<Long, CelReference> references;
545+
private final Map<ImmutableList<?>, Integer> keyToVn;
546+
private final IdentityHashMap<CelMutableExpr, Integer> exprToVn;
547+
548+
private int vnOf(CelMutableExpr expr) {
549+
return Verify.verifyNotNull(exprToVn.get(expr));
550+
}
551+
552+
private int computeVn(CelMutableExpr expr) {
553+
ImmutableList<?> key = buildKey(expr);
554+
int vn = keyToVn.computeIfAbsent(key, k -> keyToVn.size());
555+
exprToVn.put(expr, vn);
556+
return vn;
557+
}
500558

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

504650
/** 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)