From 11c17f51c7b5b72db7973e131b0c24f2cfc8e11b Mon Sep 17 00:00:00 2001 From: Sean Huh Date: Fri, 7 Aug 2026 17:30:00 -0700 Subject: [PATCH] Internal Changes PiperOrigin-RevId: 961199488 --- .../java/dev/cel/optimizer/AstMutator.java | 187 ++++++++++++++++++ .../dev/cel/optimizer/optimizers/BUILD.bazel | 1 + .../optimizers/ConstantFoldingOptimizer.java | 174 +++++++--------- .../optimizers/InliningOptimizer.java | 46 ++--- .../optimizers/SubexpressionOptimizer.java | 29 ++- .../dev/cel/optimizer/AstMutatorTest.java | 181 ++++++++++++++++- .../test/java/dev/cel/optimizer/BUILD.bazel | 1 + .../ConstantFoldingOptimizerTest.java | 4 +- 8 files changed, 472 insertions(+), 151 deletions(-) diff --git a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java index 0f428f75a..e3647de3a 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java +++ b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java @@ -18,11 +18,13 @@ import static java.lang.Math.max; import static java.util.stream.Collectors.toCollection; +import com.google.auto.value.AutoOneOf; import com.google.auto.value.AutoValue; import com.google.common.base.Preconditions; import com.google.common.base.Strings; import com.google.common.collect.HashBasedTable; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Streams; import com.google.common.collect.Table; import com.google.errorprone.annotations.Immutable; import dev.cel.common.CelAbstractSyntaxTree; @@ -44,11 +46,13 @@ import java.util.Arrays; import java.util.Collection; import java.util.HashMap; +import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map.Entry; import java.util.NoSuchElementException; import java.util.Optional; +import java.util.function.Function; import java.util.function.Predicate; import java.util.stream.Collectors; @@ -552,6 +556,140 @@ public CelMutableAst replaceSubtree( return CelMutableAst.of(mutatedRoot, newAstSource); } + /** + * Replaces a subtree in the given AST with the specified {@link SubtreeReplacement}. + * + *

This operation is intended for AST optimization purposes. + * + *

This is a very dangerous operation. Callers must re-typecheck the mutated AST and + * additionally verify that the resulting AST is semantically valid. + * + *

All expression IDs will be renumbered in a stable manner to ensure there's no ID collision + * between the nodes. The renumbering occurs even if the subtree was not replaced. + * + * @param ast Original AST to mutate. + * @param replacement Subtree replacement containing the target node ID and the new expression or + * AST. + */ + public CelMutableAst replaceSubtree(CelMutableAst ast, SubtreeReplacement replacement) { + Preconditions.checkNotNull(ast); + Preconditions.checkNotNull(replacement); + switch (replacement.replacement().kind()) { + case EXPR: + return replaceSubtree(ast, replacement.replacement().expr(), replacement.exprIdToReplace()); + case AST: + return replaceSubtree(ast, replacement.replacement().ast(), replacement.exprIdToReplace()); + } + throw new IllegalArgumentException( + "Unsupported replacement kind: " + replacement.replacement().kind()); + } + + /** + * Repeatedly applies AST mutations using the provided AST-level rewriter until no further + * replacements match (fixed point reached) or the mutator's iteration limit is exhausted. + * + *

This operation is intended for AST optimization purposes. + * + *

This is a very dangerous operation. Callers must re-typecheck the mutated AST and + * additionally verify that the resulting AST is semantically valid. + * + *

All expression IDs will be renumbered in a stable manner to ensure there's no ID collision + * between the nodes. + * + * @param ast Initial mutable AST to mutate. + * @param astRewriter Function returning a {@link SubtreeReplacement} or {@code Optional.empty()} + * when no further rewrites are possible. + * @return Mutated {@link CelMutableAst} at fixed point. + * @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}. + */ + public CelMutableAst mutateUntilFixedPoint( + CelMutableAst ast, + Function> astRewriter) { + Preconditions.checkNotNull(ast); + Preconditions.checkNotNull(astRewriter); + CelMutableAst mutableAst = ast; + for (long i = 0; i < iterationLimit; i++) { + CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(mutableAst); + Optional replacement = astRewriter.apply(navAst); + if (!replacement.isPresent()) { + return mutableAst; + } + mutableAst = replaceSubtree(mutableAst, replacement.get()); + } + throw new IllegalStateException("Max iteration count reached."); + } + + /** + * Traverses nodes using the specified {@link TraversalOrder} and repeatedly rewrites matching + * subtrees until a fixed point is reached. + * + *

This operation is intended for AST optimization purposes. + * + *

This is a very dangerous operation. Callers must re-typecheck the mutated AST and + * additionally verify that the resulting AST is semantically valid. + * + *

All expression IDs will be renumbered in a stable manner to ensure there's no ID collision + * between the nodes. + * + * @param ast Initial mutable AST to mutate. + * @param traversalOrder Order in which nodes are visited per iteration pass. + * @param nodeRewriter Function returning a {@link SubtreeReplacement} or {@code + * Optional.empty()}. + * @return Mutated {@link CelMutableAst} at fixed point. + * @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}. + */ + public CelMutableAst mutateUntilFixedPoint( + CelMutableAst ast, + TraversalOrder traversalOrder, + Function> nodeRewriter) { + Preconditions.checkNotNull(traversalOrder); + Preconditions.checkNotNull(nodeRewriter); + return mutateUntilFixedPoint( + ast, + navAst -> + navAst + .getRoot() + .allNodes(traversalOrder) + .flatMap(node -> Streams.stream(nodeRewriter.apply(node))) + .findFirst()); + } + + /** + * Traverses nodes using the specified {@link TraversalOrder}, applies the node matcher, and + * substitutes matching nodes with the returned replacement expression (targeting {@code + * node.id()}). + * + *

This operation is intended for AST optimization purposes. + * + *

This is a very dangerous operation. Callers must re-typecheck the mutated AST and + * additionally verify that the resulting AST is semantically valid. + * + *

All expression IDs will be renumbered in a stable manner to ensure there's no ID collision + * between the nodes. + * + * @param ast Initial mutable AST to mutate. + * @param traversalOrder Order in which nodes are visited per iteration pass. + * @param nodeMatcher Predicate to filter candidate nodes. + * @param nodeRewriter Function producing the new {@link CelMutableExpr} for matched nodes. + * @return Mutated {@link CelMutableAst} at fixed point. + * @throws IllegalStateException If the iteration count exceeds {@code iterationLimit}. + */ + public CelMutableAst mutateUntilFixedPoint( + CelMutableAst ast, + TraversalOrder traversalOrder, + Predicate nodeMatcher, + Function> nodeRewriter) { + Preconditions.checkNotNull(nodeMatcher); + Preconditions.checkNotNull(nodeRewriter); + return mutateUntilFixedPoint( + ast, + traversalOrder, + node -> + nodeMatcher.test(node) + ? nodeRewriter.apply(node).map(newExpr -> SubtreeReplacement.of(node.id(), newExpr)) + : Optional.empty()); + } + private CelMutableExpr mangleIdentsInComprehensionExpr( CelMutableExpr root, CelMutableExpr comprehensionExpr, @@ -983,4 +1121,53 @@ private static MangledComprehensionName of( iterVarName, iterVar2Name, resultName); } } + + /** + * Represents a planned subtree replacement containing the target node ID to replace and either a + * {@link CelMutableExpr} or {@link CelMutableAst}. + */ + @AutoValue + public abstract static class SubtreeReplacement { + + public abstract long exprIdToReplace(); + + public abstract Replacement replacement(); + + public static SubtreeReplacement of(long exprIdToReplace, CelMutableExpr replacementExpr) { + return new AutoValue_AstMutator_SubtreeReplacement( + exprIdToReplace, Replacement.ofExpr(replacementExpr)); + } + + public static SubtreeReplacement of(long exprIdToReplace, CelMutableAst replacementAst) { + return new AutoValue_AstMutator_SubtreeReplacement( + exprIdToReplace, Replacement.ofAst(replacementAst)); + } + + /** Discriminated union of either a {@link CelMutableExpr} or a {@link CelMutableAst}. */ + @AutoOneOf(Replacement.Kind.class) + public abstract static class Replacement { + + public abstract CelMutableExpr expr(); + + public abstract CelMutableAst ast(); + + public abstract Kind kind(); + + public static Replacement ofExpr(CelMutableExpr expr) { + return AutoOneOf_AstMutator_SubtreeReplacement_Replacement.expr( + Preconditions.checkNotNull(expr)); + } + + public static Replacement ofAst(CelMutableAst ast) { + return AutoOneOf_AstMutator_SubtreeReplacement_Replacement.ast( + Preconditions.checkNotNull(ast)); + } + + /** Kind of {@link Replacement}. */ + public enum Kind { + EXPR, + AST + } + } + } } diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel index 0c4b78826..1012b19c2 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/BUILD.bazel @@ -96,6 +96,7 @@ java_library( "//common:operator", "//common/ast", "//common/ast:mutable_expr", + "//common/navigation:common", "//common/navigation:expr_util", "//common/navigation:mutable_navigation", "//common/types", diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java index 266059426..5aed8acd6 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizer.java @@ -55,6 +55,7 @@ import dev.cel.common.values.StructValue; import dev.cel.extensions.CelOptionalLibrary.Function; import dev.cel.optimizer.AstMutator; +import dev.cel.optimizer.AstMutator.SubtreeReplacement; import dev.cel.optimizer.CelAstOptimizer; import dev.cel.optimizer.CelOptimizationException; import dev.cel.runtime.CelAttributePattern; @@ -134,7 +135,15 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) Cel optimizerEnv = builder.setResultType(SimpleType.DYN).build(); CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + ImmutableMap identTypes = precomputeIdentTypes(mutableAst); + mutableAst = foldConstants(optimizerEnv, valueProvider, identTypes, mutableAst); + mutableAst = pruneOptionalElements(mutableAst); + + return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst()); + } + + private static ImmutableMap precomputeIdentTypes(CelMutableAst mutableAst) { // HACK: The AstMutator strips type metadata during intermediate folds due to ID renumbering. // We pre-compute identifier types from the unmutated AST to safely evaluate boolean conditions // later. @@ -151,58 +160,47 @@ public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) Optional type = mutableAst.getType(node.id()); type.ifPresent(celType -> mutableIdentTypes.put(node.expr().ident().name(), celType)); } - ImmutableMap identTypes = ImmutableMap.copyOf(mutableIdentTypes); + return ImmutableMap.copyOf(mutableIdentTypes); + } - int iterCount = 0; - boolean continueFolding = true; - while (continueFolding) { - if (iterCount >= constantFoldingOptions.maxIterationLimit()) { - throw new IllegalStateException("Max iteration count reached."); + private CelMutableAst foldConstants( + Cel optimizerEnv, + @Nullable CelValueProvider valueProvider, + ImmutableMap identTypes, + CelMutableAst mutableAst) + throws CelOptimizationException { + for (int iterCount = 0; iterCount < constantFoldingOptions.maxIterationLimit(); iterCount++) { + Optional replacement = + findNextFoldableSubtree(optimizerEnv, valueProvider, identTypes, mutableAst); + if (!replacement.isPresent()) { + return mutableAst; } - iterCount++; - continueFolding = false; - ImmutableList foldableExprs = - CelNavigableMutableAst.fromAst(mutableAst) - .getRoot() - .allNodes(TraversalOrder.PRE_ORDER) - .filter(this::canFold) - .collect(toImmutableList()); - for (CelNavigableMutableExpr foldableExpr : foldableExprs) { - iterCount++; - - Optional mutatedResult; - // Attempt to prune if it is a non-strict call - mutatedResult = maybePruneBranches(mutableAst, identTypes, foldableExpr.expr()); - if (!mutatedResult.isPresent()) { - // Evaluate the call then fold - try { - mutatedResult = maybeFold(optimizerEnv, valueProvider, mutableAst, foldableExpr); - } catch (CelEvaluationException e) { - throw new CelOptimizationException( - "Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(), - e); - } - } - - if (!mutatedResult.isPresent()) { - // Skip this expr. It's neither prune-able nor foldable. - continue; - } + mutableAst = astMutator.replaceSubtree(mutableAst, replacement.get()); + } + throw new IllegalStateException("Max iteration count reached."); + } - continueFolding = true; - mutableAst = mutatedResult.get(); - // Break the loop because we mutated the AST. Since we traverse in PRE_ORDER (top-down), - // mutating a parent node means its children are now obsolete or folded. - // We restart the traversal to gather a fresh list of foldable expressions. - break; + private Optional findNextFoldableSubtree( + Cel optimizerEnv, + @Nullable CelValueProvider valueProvider, + ImmutableMap identTypes, + CelMutableAst mutableAst) + throws CelOptimizationException { + CelNavigableMutableAst navAst = CelNavigableMutableAst.fromAst(mutableAst); + Iterator foldableExprs = + navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).filter(this::canFold).iterator(); + while (foldableExprs.hasNext()) { + CelNavigableMutableExpr foldableExpr = foldableExprs.next(); + Optional pruned = maybePruneBranches(identTypes, foldableExpr.expr()); + if (pruned.isPresent()) { + return Optional.of(SubtreeReplacement.of(foldableExpr.id(), pruned.get())); + } + Optional folded = maybeFold(optimizerEnv, valueProvider, foldableExpr); + if (folded.isPresent()) { + return Optional.of(SubtreeReplacement.of(foldableExpr.id(), folded.get())); } } - - // If the output is a list, map, or struct which contains optional entries, then prune it - // to make sure that the optionals, if resolved, do not surface in the output literal. - mutableAst = pruneOptionalElements(mutableAst); - - return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst()); + return Optional.empty(); } private boolean canFold(CelNavigableMutableExpr navigableExpr) { @@ -317,12 +315,9 @@ private static boolean isNestedComprehension(CelNavigableMutableExpr expr) { return false; } - private Optional maybeFold( - Cel cel, - CelValueProvider valueProvider, - CelMutableAst mutableAst, - CelNavigableMutableExpr node) - throws CelOptimizationException, CelEvaluationException { + private Optional maybeFold( + Cel cel, @Nullable CelValueProvider valueProvider, CelNavigableMutableExpr node) + throws CelOptimizationException { if (!node.getKind().equals(Kind.COMPREHENSION) && CelNavigableExprUtil.hasComprehensionVariable(node)) { return Optional.empty(); @@ -330,7 +325,7 @@ private Optional maybeFold( Object result; try { result = evaluateExpr(cel, node); - } catch (CelValidationException e) { + } catch (CelEvaluationException | CelValidationException e) { throw new CelOptimizationException( "Constant folding failure. Failed to evaluate subtree due to: " + e.getMessage(), e); } @@ -340,17 +335,10 @@ private Optional maybeFold( // ex2: optional.ofNonZeroValue(5) -> optional.of(5) if (result instanceof Optional) { Optional optResult = ((Optional) result); - return maybeRewriteOptional( - cel.getTypeProvider(), valueProvider, optResult, mutableAst, node.expr()); + return maybeRewriteOptional(cel.getTypeProvider(), valueProvider, optResult, node.expr()); } - CelMutableExpr adaptedResult = - maybeAdaptEvaluatedResult(cel.getTypeProvider(), valueProvider, result).orElse(null); - if (adaptedResult == null) { - return Optional.empty(); - } - - return Optional.of(astMutator.replaceSubtree(mutableAst, adaptedResult, node.id())); + return maybeAdaptEvaluatedResult(cel.getTypeProvider(), valueProvider, result); } private Optional maybeAdaptEvaluatedResult( @@ -440,11 +428,10 @@ private Optional maybeAdaptEvaluatedResult( return Optional.empty(); } - private Optional maybeRewriteOptional( + private Optional maybeRewriteOptional( CelTypeProvider typeProvider, CelValueProvider valueProvider, Optional optResult, - CelMutableAst mutableAst, CelMutableExpr expr) { Object unwrappedResult = optResult.orElse(null); if (unwrappedResult == null) { @@ -454,7 +441,7 @@ private Optional maybeRewriteOptional( // An empty optional value was encountered. Rewrite the tree with optional.none call. // This is to account for other optional functions returning an empty optional value // e.g: optional.ofNonZeroValue(0) - return Optional.of(astMutator.replaceSubtree(mutableAst, newOptionalNoneExpr(), expr.id())); + return Optional.of(newOptionalNoneExpr()); } if (isCallToFunction(expr, Function.OPTIONAL_OF.getFunction())) { @@ -472,7 +459,7 @@ private Optional maybeRewriteOptional( CelMutableExpr.ofCall( CelMutableCall.create(Function.OPTIONAL_OF.getFunction(), adaptedResult)); - return Optional.of(astMutator.replaceSubtree(mutableAst, newOptionalOfCall, expr.id())); + return Optional.of(newOptionalOfCall); } private static boolean isCallToFunction(CelMutableExpr expr, String functionName) { @@ -480,8 +467,8 @@ private static boolean isCallToFunction(CelMutableExpr expr, String functionName } /** Inspects the non-strict calls to determine whether a branch can be removed. */ - private Optional maybePruneBranches( - CelMutableAst mutableAst, Map identTypes, CelMutableExpr expr) { + private Optional maybePruneBranches( + Map identTypes, CelMutableExpr expr) { if (!expr.getKind().equals(Kind.CALL)) { return Optional.empty(); } @@ -490,7 +477,7 @@ private Optional maybePruneBranches( String function = call.function(); if (function.equals(Operator.LOGICAL_AND.getFunction()) || function.equals(Operator.LOGICAL_OR.getFunction())) { - return maybeShortCircuitCall(mutableAst, identTypes, expr); + return maybeShortCircuitCall(identTypes, expr); } else if (function.equals(Operator.CONDITIONAL.getFunction())) { CelMutableExpr cond = call.args().get(0); CelMutableExpr truthy = call.args().get(1); @@ -501,7 +488,7 @@ private Optional maybePruneBranches( } CelMutableExpr result = cond.constant().booleanValue() ? truthy : falsy; - return Optional.of(astMutator.replaceSubtree(mutableAst, result, expr.id())); + return Optional.of(result); } else if (function.equals(Operator.IN.getFunction())) { CelMutableExpr callArg = call.args().get(1); if (!callArg.getKind().equals(Kind.LIST)) { @@ -510,9 +497,7 @@ private Optional maybePruneBranches( CelMutableList haystack = callArg.list(); if (haystack.elements().isEmpty()) { - return Optional.of( - astMutator.replaceSubtree( - mutableAst, CelMutableExpr.ofConstant(CelConstant.ofValue(false)), expr.id())); + return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(false))); } CelMutableExpr needle = call.args().get(0); @@ -529,24 +514,13 @@ private Optional maybePruneBranches( && Double.isNaN(needle.constant().doubleValue())) { continue; } - return Optional.of( - astMutator.replaceSubtree( - mutableAst.expr(), - CelMutableExpr.ofConstant(CelConstant.ofValue(true)), - expr.id())); + return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(true))); } - CelType needleType = - mutableAst - .getType(needle.id()) - .orElseGet(() -> identTypes.get(needle.ident().name())); + CelType needleType = identTypes.get(needle.ident().name()); if (needleType != null && isSafeForExactEquality(needleType)) { - return Optional.of( - astMutator.replaceSubtree( - mutableAst.expr(), - CelMutableExpr.ofConstant(CelConstant.ofValue(true)), - expr.id())); + return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(true))); } } } @@ -566,7 +540,7 @@ private Optional maybePruneBranches( return Optional.empty(); } else if (lhsIsBooleanConstant && (!constantFoldingOptions.enableSafeLogicalOptimization() - || evaluatesToBoolean(mutableAst, identTypes, rhs))) { + || evaluatesToBoolean(identTypes, rhs))) { boolean cond = invertCondition != lhs.constant().booleanValue(); replacementExpr = Optional.of( @@ -576,7 +550,7 @@ private Optional maybePruneBranches( CelMutableCall.create(Operator.LOGICAL_NOT.getFunction(), rhs))); } else if (rhsIsBooleanConstant && (!constantFoldingOptions.enableSafeLogicalOptimization() - || evaluatesToBoolean(mutableAst, identTypes, lhs))) { + || evaluatesToBoolean(identTypes, lhs))) { boolean cond = invertCondition != rhs.constant().booleanValue(); replacementExpr = Optional.of( @@ -586,14 +560,14 @@ private Optional maybePruneBranches( CelMutableCall.create(Operator.LOGICAL_NOT.getFunction(), lhs))); } - return replacementExpr.map(node -> astMutator.replaceSubtree(mutableAst, node, expr.id())); + return replacementExpr; } return Optional.empty(); } - private Optional maybeShortCircuitCall( - CelMutableAst mutableAst, Map identTypes, CelMutableExpr expr) { + private Optional maybeShortCircuitCall( + Map identTypes, CelMutableExpr expr) { CelMutableCall call = expr.call(); boolean shortCircuit = false; boolean skip = true; @@ -613,7 +587,7 @@ private Optional maybeShortCircuitCall( } if (arg.constant().booleanValue() == shortCircuit) { - return Optional.of(astMutator.replaceSubtree(mutableAst, arg, expr.id())); + return Optional.of(arg); } } @@ -621,13 +595,13 @@ private Optional maybeShortCircuitCall( if (newArgs.isEmpty()) { CelMutableExpr shortCircuitTarget = call.args().get(0); // either args(0) or args(1) would work here - return Optional.of(astMutator.replaceSubtree(mutableAst, shortCircuitTarget, expr.id())); + return Optional.of(shortCircuitTarget); } if (newArgs.size() == 1) { CelMutableExpr remainingArg = newArgs.get(0); if (!constantFoldingOptions.enableSafeLogicalOptimization() - || evaluatesToBoolean(mutableAst, identTypes, remainingArg)) { - return Optional.of(astMutator.replaceSubtree(mutableAst, remainingArg, expr.id())); + || evaluatesToBoolean(identTypes, remainingArg)) { + return Optional.of(remainingArg); } return Optional.empty(); } @@ -637,8 +611,8 @@ private Optional maybeShortCircuitCall( "Folding variadic logical operator is not supported yet."); } - private boolean evaluatesToBoolean( - CelMutableAst mutableAst, Map identTypes, CelMutableExpr expr) { + private static boolean evaluatesToBoolean( + Map identTypes, CelMutableExpr expr) { if (isExprConstantOfKind(expr, CelConstant.Kind.BOOLEAN_VALUE)) { return true; } @@ -654,7 +628,7 @@ private boolean evaluatesToBoolean( if (expr.getKind().equals(Kind.IDENT)) { return Objects.equals(identTypes.get(expr.ident().name()), SimpleType.BOOL); } - return mutableAst.getType(expr.id()).map(SimpleType.BOOL::equals).orElse(false); + return false; } private boolean isFoldedAggregateLiteral(CelMutableExpr expr) { diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java index 147673e47..61fd19347 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/InliningOptimizer.java @@ -14,8 +14,6 @@ package dev.cel.optimizer.optimizers; -import static com.google.common.collect.ImmutableList.toImmutableList; - import com.google.auto.value.AutoValue; import com.google.common.collect.ImmutableList; import com.google.common.primitives.UnsignedLong; @@ -29,13 +27,14 @@ import dev.cel.common.ast.CelMutableExpr.CelMutableCall; import dev.cel.common.ast.CelMutableExpr.CelMutableStruct; import dev.cel.common.navigation.CelNavigableExprUtil; -import dev.cel.common.navigation.CelNavigableMutableAst; import dev.cel.common.navigation.CelNavigableMutableExpr; +import dev.cel.common.navigation.TraversalOrder; import dev.cel.common.types.CelKind; import dev.cel.common.types.CelType; import dev.cel.common.types.SimpleType; import dev.cel.common.values.NullValue; import dev.cel.optimizer.AstMutator; +import dev.cel.optimizer.AstMutator.SubtreeReplacement; import dev.cel.optimizer.CelAstOptimizer; import java.util.ArrayList; import java.util.List; @@ -105,28 +104,25 @@ public static InliningOptimizer newInstance( public OptimizationResult optimize(CelAbstractSyntaxTree ast, Cel cel) { CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); for (InlineVariable inlineVariable : inlineVariables) { - ImmutableList inlinableExprs = - CelNavigableMutableAst.fromAst(mutableAst) - .getRoot() - .allNodes() - .filter(node -> canInline(node, inlineVariable.name())) - .collect(toImmutableList()); - - for (CelNavigableMutableExpr inlinableExpr : inlinableExprs) { - CelMutableAst inlineVariableAst = CelMutableAst.fromCelAst(inlineVariable.ast()); - CelMutableExpr replacementExpr = inlineVariableAst.expr(); - - if (inlinableExpr.getKind().equals(Kind.SELECT) - && inlinableExpr.expr().select().testOnly()) { - replacementExpr = rewritePresenceExpr(inlineVariable, replacementExpr); - } - - mutableAst = - astMutator.replaceSubtree( - mutableAst, - CelMutableAst.of(replacementExpr, inlineVariableAst.source()), - inlinableExpr.id()); - } + mutableAst = + astMutator.mutateUntilFixedPoint( + mutableAst, + TraversalOrder.POST_ORDER, + node -> { + if (!canInline(node, inlineVariable.name())) { + return Optional.empty(); + } + CelMutableAst inlineVariableAst = CelMutableAst.fromCelAst(inlineVariable.ast()); + CelMutableExpr replacementExpr = inlineVariableAst.expr(); + + if (node.getKind().equals(Kind.SELECT) && node.expr().select().testOnly()) { + replacementExpr = rewritePresenceExpr(inlineVariable, replacementExpr); + } + + return Optional.of( + SubtreeReplacement.of( + node.id(), CelMutableAst.of(replacementExpr, inlineVariableAst.source()))); + }); } return OptimizationResult.create(astMutator.renumberIdsConsecutively(mutableAst).toParsedAst()); diff --git a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java index 5eebb1c54..6a9860750 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java @@ -65,6 +65,7 @@ import java.util.Comparator; import java.util.HashSet; import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.stream.Stream; @@ -161,27 +162,23 @@ private OptimizationResult optimizeUsingCelBlock(CelAbstractSyntaxTree ast, Cel break; } + CelMutableExpr targetCseShape = normalizeForEquality(cseCandidates.get(0)); subexpressions.add(cseCandidates.get(0)); String blockIdentifier = BLOCK_INDEX_PREFIX + blockIdentifierIndex++; // Replace all CSE candidates with new block index identifier - for (CelMutableExpr cseCandidate : cseCandidates) { - iterCount++; - - astToModify = - astMutator.replaceSubtree( - navAst, - CelNavigableMutableAst.fromAst( - CelMutableAst.of( - CelMutableExpr.ofIdent(blockIdentifier), navAst.getAst().source())), - cseCandidate.id()); - - // Retain the existing macro calls in case if the block identifiers are replacing a subtree - // that contains a comprehension. - sourceToModify.addAllMacroCalls(astToModify.source().getMacroCalls()); - astToModify = CelMutableAst.of(astToModify.expr(), sourceToModify); - } + astToModify = + astMutator.mutateUntilFixedPoint( + astToModify, + TraversalOrder.POST_ORDER, + node -> normalizeForEquality(node.expr()).equals(targetCseShape), + node -> Optional.of(CelMutableExpr.ofIdent(blockIdentifier))); + + // Retain the existing macro calls in case if the block identifiers are replacing a subtree + // that contains a comprehension. + sourceToModify.addAllMacroCalls(astToModify.source().getMacroCalls()); + astToModify = CelMutableAst.of(astToModify.expr(), sourceToModify); } if (iterCount >= cseOptions.iterationLimit()) { diff --git a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java index fa896ebca..f0c3a7045 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java @@ -41,24 +41,31 @@ import dev.cel.common.ast.CelMutableExprConverter; import dev.cel.common.navigation.CelNavigableAst; import dev.cel.common.navigation.CelNavigableExpr; +import dev.cel.common.navigation.TraversalOrder; import dev.cel.common.types.SimpleType; import dev.cel.common.types.StructTypeReference; import dev.cel.expr.conformance.proto3.TestAllTypes; import dev.cel.extensions.CelExtensions; import dev.cel.extensions.CelOptionalLibrary; +import dev.cel.optimizer.AstMutator.SubtreeReplacement; import dev.cel.parser.CelStandardMacro; import dev.cel.parser.CelUnparser; import dev.cel.parser.CelUnparserFactory; import java.util.List; +import java.util.Optional; import org.junit.Test; import org.junit.runner.RunWith; @RunWith(TestParameterInjector.class) public class AstMutatorTest { private static final Cel CEL = - CelFactory.standardCelBuilder() + CelFactory.plannerCelBuilder() .setStandardMacros(CelStandardMacro.STANDARD_MACROS) - .setOptions(CelOptions.current().populateMacroCalls(true).build()) + .setOptions( + CelOptions.current() + .populateMacroCalls(true) + .enableHeterogeneousNumericComparisons(true) + .build()) .addMessageTypes(TestAllTypes.getDescriptor()) .addCompilerLibraries( CelOptionalLibrary.INSTANCE, CelExtensions.bindings(), CelExtensions.comprehensions()) @@ -66,6 +73,7 @@ public class AstMutatorTest { .setContainer(CelContainer.ofName("cel.expr.conformance.proto3")) .addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName())) .addVar("x", SimpleType.INT) + .addVar("b", SimpleType.BOOL) .build(); private static final CelUnparser CEL_UNPARSER = CelUnparserFactory.newUnparser(); @@ -440,7 +448,7 @@ public void memberCallExpr_replaceLeafTarget() throws Exception { // 10 [1] func [4] // 4 [3] 5 [5] Cel cel = - CelFactory.standardCelBuilder() + CelFactory.plannerCelBuilder() .addFunctionDeclarations( CelFunctionDecl.newFunctionDeclaration( "func", @@ -463,7 +471,7 @@ public void memberCallExpr_replaceLeafArgument() throws Exception { // 10 [1] func [4] // 4 [3] 5 [5] Cel cel = - CelFactory.standardCelBuilder() + CelFactory.plannerCelBuilder() .addFunctionDeclarations( CelFunctionDecl.newFunctionDeclaration( "func", @@ -486,7 +494,7 @@ public void memberCallExpr_replaceMiddleBranchTarget() throws Exception { // 10 [1] func [4] // 4 [3] 5 [5] Cel cel = - CelFactory.standardCelBuilder() + CelFactory.plannerCelBuilder() .addFunctionDeclarations( CelFunctionDecl.newFunctionDeclaration( "func", @@ -509,7 +517,7 @@ public void memberCallExpr_replaceMiddleBranchArgument() throws Exception { // 10 [1] func [4] // 4 [3] 5 [5] Cel cel = - CelFactory.standardCelBuilder() + CelFactory.plannerCelBuilder() .addFunctionDeclarations( CelFunctionDecl.newFunctionDeclaration( "func", @@ -864,9 +872,13 @@ public void mangleComprehensionVariable_adjacentMacros_differentIterVarTypes() t public void mangleComprehensionVariable_macroSourceDisabled_macroCallMapIsEmpty() throws Exception { Cel cel = - CelFactory.standardCelBuilder() + CelFactory.plannerCelBuilder() .setStandardMacros(CelStandardMacro.STANDARD_MACROS) - .setOptions(CelOptions.current().populateMacroCalls(false).build()) + .setOptions( + CelOptions.current() + .populateMacroCalls(false) + .enableHeterogeneousNumericComparisons(true) + .build()) .build(); CelAbstractSyntaxTree ast = cel.compile("[false].exists(i, i)").getAst(); @@ -1044,6 +1056,159 @@ public void newGlobalCallAst_success() throws Exception { .isEqualTo("func([1].exists(x, x >= 1), \"hello\")"); } + @Test + public void replaceSubtree_withSubtreeReplacement_expr() throws Exception { + CelAbstractSyntaxTree ast = CEL.compile("1 + 2").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + SubtreeReplacement replacement = + SubtreeReplacement.of(1, CelMutableExpr.ofConstant(CelConstant.ofValue(10))); + + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, replacement); + + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10 + 2"); + } + + @Test + public void replaceSubtree_withSubtreeReplacement_ast() throws Exception { + CelAbstractSyntaxTree ast = CEL.compile("true && false").getAst(); + CelAbstractSyntaxTree macroAst = CEL.compile("[1].exists(x, x > 0)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + SubtreeReplacement replacement = SubtreeReplacement.of(3, CelMutableAst.fromCelAst(macroAst)); + + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, replacement); + + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())) + .isEqualTo("true && [1].exists(x, x > 0)"); + assertThat(result.source().getMacroCalls()).hasSize(1); + } + + @Test + public void mutateUntilFixedPoint_astRewriter_success() throws Exception { + // Repeatedly simplifies addition with 0: "1 + 0 + 0" -> "1" + CelAbstractSyntaxTree ast = CEL.compile("1 + 0 + 0").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + + CelMutableAst result = + AST_MUTATOR.mutateUntilFixedPoint( + mutableAst, + navAst -> + navAst + .getRoot() + .allNodes() + .filter( + node -> + node.getKind().equals(Kind.CALL) + && node.expr().call().function().equals("_+_")) + .filter( + node -> { + List args = node.expr().call().args(); + return args.get(1).getKind().equals(Kind.CONSTANT) + && args.get(1).constant().int64Value() == 0; + }) + .map(node -> SubtreeReplacement.of(node.id(), node.expr().call().args().get(0))) + .findFirst()); + + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("1"); + } + + @Test + public void mutateUntilFixedPoint_nodeRewriter_success() throws Exception { + // Rewrites nested calls: "func(func(1))" -> "1" + Cel cel = + CelFactory.plannerCelBuilder() + .addFunctionDeclarations( + CelFunctionDecl.newFunctionDeclaration( + "func", + CelOverloadDecl.newGlobalOverload( + "func_overload", SimpleType.INT, SimpleType.INT))) + .build(); + CelAbstractSyntaxTree ast = cel.compile("func(func(10))").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + + CelMutableAst result = + AST_MUTATOR.mutateUntilFixedPoint( + mutableAst, + TraversalOrder.POST_ORDER, + node -> { + if (node.getKind().equals(Kind.CALL) + && node.expr().call().function().equals("func")) { + return Optional.of( + SubtreeReplacement.of(node.id(), node.expr().call().args().get(0))); + } + return Optional.empty(); + }); + + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10"); + } + + @Test + public void mutateUntilFixedPoint_matcherAndNodeRewriter_success() throws Exception { + // Replaces all variables named 'x' with constant 5 in "x + x + x" -> "5 + 5 + 5" + CelAbstractSyntaxTree ast = CEL.compile("x + x + x").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + + CelMutableAst result = + AST_MUTATOR.mutateUntilFixedPoint( + mutableAst, + TraversalOrder.POST_ORDER, + node -> node.getKind().equals(Kind.IDENT) && node.expr().ident().name().equals("x"), + node -> Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(5)))); + + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("5 + 5 + 5"); + } + + @Test + public void mutateUntilFixedPoint_withReplacementAst_preservesMacroSource() throws Exception { + // Replaces identifier 'b' with macro AST in "b && true" + CelAbstractSyntaxTree ast = CEL.compile("b && true").getAst(); + CelAbstractSyntaxTree macroAst = CEL.compile("[1].exists(i, i > 0)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + + CelMutableAst result = + AST_MUTATOR.mutateUntilFixedPoint( + mutableAst, + navAst -> + navAst + .getRoot() + .allNodes() + .filter( + node -> + node.getKind().equals(Kind.IDENT) + && node.expr().ident().name().equals("b")) + .map( + node -> + SubtreeReplacement.of(node.id(), CelMutableAst.fromCelAst(macroAst))) + .findFirst()); + + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())) + .isEqualTo("[1].exists(i, i > 0) && true"); + assertThat(result.source().getMacroCalls()).hasSize(1); + assertThat(CEL.createProgram(CEL.check(result.toParsedAst()).getAst()).eval()).isEqualTo(true); + } + + @Test + public void mutateUntilFixedPoint_exceedsIterationLimit_throws() throws Exception { + // Circular rewrite rule that alternates between 1 and 2 + CelAbstractSyntaxTree ast = CEL.compile("1").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + + IllegalStateException e = + assertThrows( + IllegalStateException.class, + () -> + AST_MUTATOR.mutateUntilFixedPoint( + mutableAst, + TraversalOrder.POST_ORDER, + node -> node.getKind().equals(Kind.CONSTANT), + node -> { + long currentVal = node.expr().constant().int64Value(); + long nextVal = currentVal == 1L ? 2L : 1L; + return Optional.of(CelMutableExpr.ofConstant(CelConstant.ofValue(nextVal))); + })); + + assertThat(e).hasMessageThat().isEqualTo("Max iteration count reached."); + } + @Test public void newMemberCallAst_success() throws Exception { CelMutableAst targetAst = CelMutableAst.fromCelAst(CEL.compile("'hello'").getAst()); diff --git a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel index 8ea72a261..702fe23f3 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel +++ b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel @@ -21,6 +21,7 @@ java_library( "//common/ast", "//common/ast:mutable_expr", "//common/navigation", + "//common/navigation:common", "//common/types", "//compiler", "//extensions", diff --git a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java index 613b53ea3..d00b7beac 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/optimizers/ConstantFoldingOptimizerTest.java @@ -128,13 +128,13 @@ private static Cel setupEnv(CelBuilder celBuilder) { CelExtensions.comprehensions(), CelExtensions.bindings(), CelOptionalLibrary.INSTANCE, - CelExtensions.math(CEL_OPTIONS), + CelExtensions.math(), CelExtensions.strings(), CelExtensions.sets(CEL_OPTIONS), CelExtensions.encoders(CEL_OPTIONS)) .addRuntimeLibraries( CelOptionalLibrary.INSTANCE, - CelExtensions.math(CEL_OPTIONS), + CelExtensions.math(), CelExtensions.strings(), CelExtensions.sets(CEL_OPTIONS), CelExtensions.encoders(CEL_OPTIONS))