Skip to content

Commit 2f6380f

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Add helpers for performing fixed point optimization
PiperOrigin-RevId: 961182997
1 parent d86bfbe commit 2f6380f

8 files changed

Lines changed: 499 additions & 151 deletions

File tree

‎optimizer/src/main/java/dev/cel/optimizer/AstMutator.java‎

Lines changed: 209 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,17 +18,19 @@
1818
import static java.lang.Math.max;
1919
import static java.util.stream.Collectors.toCollection;
2020

21+
import com.google.auto.value.AutoOneOf;
2122
import com.google.auto.value.AutoValue;
2223
import com.google.common.base.Preconditions;
2324
import com.google.common.base.Strings;
2425
import com.google.common.collect.HashBasedTable;
2526
import com.google.common.collect.ImmutableMap;
27+
import com.google.common.collect.Streams;
2628
import com.google.common.collect.Table;
2729
import com.google.errorprone.annotations.Immutable;
2830
import dev.cel.common.CelAbstractSyntaxTree;
2931
import dev.cel.common.CelMutableAst;
3032
import dev.cel.common.CelMutableSource;
31-
import dev.cel.common.ast.CelExpr.ExprKind.Kind;
33+
import dev.cel.common.ast.CelExpr.ExprKind;
3234
import dev.cel.common.ast.CelExprIdGeneratorFactory;
3335
import dev.cel.common.ast.CelExprIdGeneratorFactory.ExprIdGenerator;
3436
import dev.cel.common.ast.CelExprIdGeneratorFactory.StableIdGenerator;
@@ -49,6 +51,7 @@
4951
import java.util.Map.Entry;
5052
import java.util.NoSuchElementException;
5153
import java.util.Optional;
54+
import java.util.function.Function;
5255
import java.util.function.Predicate;
5356
import java.util.stream.Collectors;
5457

@@ -213,7 +216,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
213216
Predicate<CelNavigableMutableExpr> comprehensionIdentifierPredicate = x -> true;
214217
comprehensionIdentifierPredicate =
215218
comprehensionIdentifierPredicate
216-
.and(node -> node.getKind().equals(Kind.COMPREHENSION))
219+
.and(node -> node.getKind().equals(ExprKind.Kind.COMPREHENSION))
217220
.and(node -> !node.expr().comprehension().iterVar().startsWith(newIterVarPrefix + ":"))
218221
.and(node -> !node.expr().comprehension().accuVar().startsWith(newAccuVarPrefix + ":"))
219222
.and(
@@ -235,7 +238,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
235238
String result = node.expr().comprehension().result().ident().name();
236239
return CelNavigableMutableExpr.fromExpr(node.expr().comprehension().loopStep())
237240
.allNodes()
238-
.filter(subNode -> subNode.getKind().equals(Kind.IDENT))
241+
.filter(subNode -> subNode.getKind().equals(ExprKind.Kind.IDENT))
239242
.map(subNode -> subNode.expr().ident())
240243
.anyMatch(
241244
ident ->
@@ -259,7 +262,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
259262
.allNodes()
260263
.filter(
261264
loopStepNode ->
262-
loopStepNode.getKind().equals(Kind.IDENT)
265+
loopStepNode.getKind().equals(ExprKind.Kind.IDENT)
263266
&& loopStepNode.expr().ident().name().equals(iterVar))
264267
.map(CelNavigableMutableExpr::id)
265268
.findAny();
@@ -269,7 +272,7 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames(
269272
.filter(
270273
loopStepNode ->
271274
!iterVar2.isEmpty()
272-
&& loopStepNode.getKind().equals(Kind.IDENT)
275+
&& loopStepNode.getKind().equals(ExprKind.Kind.IDENT)
273276
&& loopStepNode.expr().ident().name().equals(iterVar2))
274277
.map(CelNavigableMutableExpr::id)
275278
.findAny();
@@ -406,13 +409,13 @@ private static MangledComprehensionName getMangledComprehensionName(
406409
private static int countComprehensionNestingLevel(CelNavigableMutableExpr comprehensionExpr) {
407410
return comprehensionExpr
408411
.descendants()
409-
.filter(node -> node.getKind().equals(Kind.COMPREHENSION))
412+
.filter(node -> node.getKind().equals(ExprKind.Kind.COMPREHENSION))
410413
.mapToInt(
411414
node -> {
412415
int nestedLevel = 1;
413416
CelNavigableMutableExpr maybeParent = node.parent().orElse(null);
414417
while (maybeParent != null && maybeParent.id() != comprehensionExpr.id()) {
415-
if (maybeParent.getKind().equals(Kind.COMPREHENSION)) {
418+
if (maybeParent.getKind().equals(ExprKind.Kind.COMPREHENSION)) {
416419
nestedLevel++;
417420
}
418421
maybeParent = maybeParent.parent().orElse(null);
@@ -552,6 +555,151 @@ public CelMutableAst replaceSubtree(
552555
return CelMutableAst.of(mutatedRoot, newAstSource);
553556
}
554557

558+
/**
559+
* Replaces a subtree in the given AST with the specified {@link SubtreeReplacement}.
560+
*
561+
* <p>This operation is intended for AST optimization purposes.
562+
*
563+
* <p>This is a very dangerous operation. Callers must re-typecheck the mutated AST and
564+
* additionally verify that the resulting AST is semantically valid.
565+
*
566+
* <p>All expression IDs will be renumbered in a stable manner to ensure there's no ID collision
567+
* between the nodes. The renumbering occurs even if the subtree was not replaced.
568+
*
569+
* @param ast Original AST to mutate.
570+
* @param replacement Subtree replacement containing the target node ID and the new expression or
571+
* AST.
572+
*/
573+
public CelMutableAst replaceSubtree(CelMutableAst ast, SubtreeReplacement replacement) {
574+
Preconditions.checkNotNull(ast);
575+
Preconditions.checkNotNull(replacement);
576+
switch (replacement.replacement().kind()) {
577+
case EXPR:
578+
return replaceSubtree(ast, replacement.replacement().expr(), replacement.exprIdToReplace());
579+
case AST:
580+
return replaceSubtree(ast, replacement.replacement().ast(), replacement.exprIdToReplace());
581+
}
582+
throw new IllegalArgumentException(
583+
"Unsupported replacement kind: " + replacement.replacement().kind());
584+
}
585+
586+
/**
587+
* Repeatedly applies AST mutations using the provided AST-level rewriter until no further
588+
* replacements match (fixed point reached) or the mutator's iteration limit is exhausted.
589+
*
590+
* <p>Per iteration pass, the AST is traversed to find and perform at most one matching subtree
591+
* substitution before restarting traversal on the newly mutated AST.
592+
*
593+
* <p>This operation is intended for AST optimization purposes.
594+
*
595+
* <p>This is a very dangerous operation. Callers must re-typecheck the mutated AST and
596+
* additionally verify that the resulting AST is semantically valid.
597+
*
598+
* <p>All expression IDs will be renumbered in a stable manner to ensure there's no ID collision
599+
* between the nodes.
600+
*
601+
* @param ast Initial mutable AST to mutate.
602+
* @param astRewriter Function returning a {@link SubtreeReplacement} or {@code Optional.empty()}
603+
* when no further rewrites are possible.
604+
* @return Mutated {@link CelMutableAst} at fixed point.
605+
* @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}.
606+
*/
607+
public CelMutableAst mutateUntilFixedPoint(
608+
CelMutableAst ast,
609+
Function<CelNavigableMutableAst, Optional<SubtreeReplacement>> astRewriter) {
610+
Preconditions.checkNotNull(ast);
611+
Preconditions.checkNotNull(astRewriter);
612+
CelMutableAst mutableAst = ast;
613+
for (long i = 0; i < iterationLimit; i++) {
614+
CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(mutableAst);
615+
Optional<SubtreeReplacement> replacement = astRewriter.apply(navAst);
616+
if (!replacement.isPresent()) {
617+
return mutableAst;
618+
}
619+
mutableAst = replaceSubtree(mutableAst, replacement.get());
620+
}
621+
throw new IllegalStateException("Max iteration count reached.");
622+
}
623+
624+
/**
625+
* Traverses nodes using the specified {@link TraversalOrder} and repeatedly rewrites matching
626+
* subtrees until a fixed point is reached.
627+
*
628+
* <p>Per iteration pass, the AST is walked in the given order to find and perform the first
629+
* matching substitution. The traversal then restarts on the freshly mutated AST until no nodes
630+
* match or the mutator's iteration limit is exhausted.
631+
*
632+
* <p>This operation is intended for AST optimization purposes.
633+
*
634+
* <p>This is a very dangerous operation. Callers must re-typecheck the mutated AST and
635+
* additionally verify that the resulting AST is semantically valid.
636+
*
637+
* <p>All expression IDs will be renumbered in a stable manner to ensure there's no ID collision
638+
* between the nodes.
639+
*
640+
* @param ast Initial mutable AST to mutate.
641+
* @param traversalOrder Order in which nodes are visited per iteration pass.
642+
* @param nodeRewriter Function returning a {@link SubtreeReplacement} or {@code
643+
* Optional.empty()}.
644+
* @return Mutated {@link CelMutableAst} at fixed point.
645+
* @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}.
646+
*/
647+
public CelMutableAst mutateUntilFixedPoint(
648+
CelMutableAst ast,
649+
TraversalOrder traversalOrder,
650+
Function<CelNavigableMutableExpr, Optional<SubtreeReplacement>> nodeRewriter) {
651+
Preconditions.checkNotNull(traversalOrder);
652+
Preconditions.checkNotNull(nodeRewriter);
653+
return mutateUntilFixedPoint(
654+
ast,
655+
navAst ->
656+
navAst
657+
.getRoot()
658+
.allNodes(traversalOrder)
659+
.flatMap(node -> Streams.stream(nodeRewriter.apply(node)))
660+
.findFirst());
661+
}
662+
663+
/**
664+
* Traverses nodes using the specified {@link TraversalOrder}, applies the node matcher, and
665+
* substitutes matching nodes with the returned replacement expression (targeting {@code
666+
* node.id()}).
667+
*
668+
* <p>Per iteration pass, the AST is walked in the given order to find and perform the first
669+
* matching substitution. The traversal then restarts on the freshly mutated AST until no nodes
670+
* match or the mutator's iteration limit is exhausted.
671+
*
672+
* <p>This operation is intended for AST optimization purposes.
673+
*
674+
* <p>This is a very dangerous operation. Callers must re-typecheck the mutated AST and
675+
* additionally verify that the resulting AST is semantically valid.
676+
*
677+
* <p>All expression IDs will be renumbered in a stable manner to ensure there's no ID collision
678+
* between the nodes.
679+
*
680+
* @param ast Initial mutable AST to mutate.
681+
* @param traversalOrder Order in which nodes are visited per iteration pass.
682+
* @param nodeMatcher Predicate to filter candidate nodes.
683+
* @param nodeRewriter Function producing the new {@link CelMutableExpr} for matched nodes.
684+
* @return Mutated {@link CelMutableAst} at fixed point.
685+
* @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}.
686+
*/
687+
public CelMutableAst mutateUntilFixedPoint(
688+
CelMutableAst ast,
689+
TraversalOrder traversalOrder,
690+
Predicate<CelNavigableMutableExpr> nodeMatcher,
691+
Function<CelNavigableMutableExpr, Optional<CelMutableExpr>> nodeRewriter) {
692+
Preconditions.checkNotNull(nodeMatcher);
693+
Preconditions.checkNotNull(nodeRewriter);
694+
return mutateUntilFixedPoint(
695+
ast,
696+
traversalOrder,
697+
node ->
698+
nodeMatcher.test(node)
699+
? nodeRewriter.apply(node).map(newExpr -> SubtreeReplacement.of(node.id(), newExpr))
700+
: Optional.empty());
701+
}
702+
555703
private CelMutableExpr mangleIdentsInComprehensionExpr(
556704
CelMutableExpr root,
557705
CelMutableExpr comprehensionExpr,
@@ -590,7 +738,7 @@ private void replaceIdentName(
590738
.map(CelNavigableMutableExpr::expr)
591739
.filter(
592740
node ->
593-
node.getKind().equals(Kind.IDENT)
741+
node.getKind().equals(ExprKind.Kind.IDENT)
594742
&& node.ident().name().equals(originalIdentName))
595743
.findAny()
596744
.orElse(null);
@@ -776,7 +924,7 @@ private CelMutableSource normalizeMacroSource(
776924
long replacedId = idGenerator.generate(exprIdToReplace);
777925
boolean isListExprBeingReplaced =
778926
allExprs.containsKey(replacedId)
779-
&& allExprs.get(replacedId).getKind().equals(Kind.LIST);
927+
&& allExprs.get(replacedId).getKind().equals(ExprKind.Kind.LIST);
780928
if (isListExprBeingReplaced) {
781929
unwrapListArgumentsInMacroCallExpr(
782930
allExprs.get(callId).comprehension(), newMacroCallExpr);
@@ -791,7 +939,7 @@ private CelMutableSource normalizeMacroSource(
791939
CelMutableExpr macroCallExpr = macroCall.getValue();
792940
CelNavigableMutableExpr.fromExpr(macroCallExpr)
793941
.allNodes()
794-
.filter(node -> node.getKind().equals(Kind.COMPREHENSION))
942+
.filter(node -> node.getKind().equals(ExprKind.Kind.COMPREHENSION))
795943
.map(CelNavigableMutableExpr::expr)
796944
.forEach(
797945
node -> {
@@ -808,7 +956,7 @@ private CelMutableSource normalizeMacroSource(
808956
// This can occur from pulling out a nested comprehension into a separate cel.block index
809957
CelNavigableMutableExpr.fromExpr(macroCallExpr)
810958
.allNodes()
811-
.filter(node -> node.getKind().equals(Kind.NOT_SET))
959+
.filter(node -> node.getKind().equals(ExprKind.Kind.NOT_SET))
812960
.map(CelNavigableMutableExpr::id)
813961
.filter(id -> !allExprs.containsKey(id))
814962
.forEach(
@@ -840,7 +988,7 @@ private CelMutableSource normalizeMacroSource(
840988
private static void unwrapListArgumentsInMacroCallExpr(
841989
CelMutableComprehension comprehension, CelMutableExpr newMacroCallExpr) {
842990
CelMutableExpr accuInit = comprehension.accuInit();
843-
if (!accuInit.getKind().equals(Kind.LIST) || !accuInit.list().elements().isEmpty()) {
991+
if (!accuInit.getKind().equals(ExprKind.Kind.LIST) || !accuInit.list().elements().isEmpty()) {
844992
// Does not contain an extraneous list.
845993
return;
846994
}
@@ -983,4 +1131,53 @@ private static MangledComprehensionName of(
9831131
iterVarName, iterVar2Name, resultName);
9841132
}
9851133
}
1134+
1135+
/**
1136+
* Represents a planned subtree replacement containing the target node ID to replace and either a
1137+
* {@link CelMutableExpr} or {@link CelMutableAst}.
1138+
*/
1139+
@AutoValue
1140+
public abstract static class SubtreeReplacement {
1141+
1142+
public abstract long exprIdToReplace();
1143+
1144+
public abstract Replacement replacement();
1145+
1146+
public static SubtreeReplacement of(long exprIdToReplace, CelMutableExpr replacementExpr) {
1147+
return new AutoValue_AstMutator_SubtreeReplacement(
1148+
exprIdToReplace, Replacement.ofExpr(replacementExpr));
1149+
}
1150+
1151+
public static SubtreeReplacement of(long exprIdToReplace, CelMutableAst replacementAst) {
1152+
return new AutoValue_AstMutator_SubtreeReplacement(
1153+
exprIdToReplace, Replacement.ofAst(replacementAst));
1154+
}
1155+
1156+
/** Discriminated union of either a {@link CelMutableExpr} or a {@link CelMutableAst}. */
1157+
@AutoOneOf(Replacement.Kind.class)
1158+
public abstract static class Replacement {
1159+
1160+
public abstract CelMutableExpr expr();
1161+
1162+
public abstract CelMutableAst ast();
1163+
1164+
public abstract Replacement.Kind kind();
1165+
1166+
public static Replacement ofExpr(CelMutableExpr expr) {
1167+
return AutoOneOf_AstMutator_SubtreeReplacement_Replacement.expr(
1168+
Preconditions.checkNotNull(expr));
1169+
}
1170+
1171+
public static Replacement ofAst(CelMutableAst ast) {
1172+
return AutoOneOf_AstMutator_SubtreeReplacement_Replacement.ast(
1173+
Preconditions.checkNotNull(ast));
1174+
}
1175+
1176+
/** Kind of {@link Replacement}. */
1177+
public enum Kind {
1178+
EXPR,
1179+
AST
1180+
}
1181+
}
1182+
}
9861183
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ java_library(
9696
"//common:operator",
9797
"//common/ast",
9898
"//common/ast:mutable_expr",
99+
"//common/navigation:common",
99100
"//common/navigation:expr_util",
100101
"//common/navigation:mutable_navigation",
101102
"//common/types",

0 commit comments

Comments
 (0)