diff --git a/common/src/main/java/dev/cel/common/CelSource.java b/common/src/main/java/dev/cel/common/CelSource.java index 6adaf877..2422f7b2 100644 --- a/common/src/main/java/dev/cel/common/CelSource.java +++ b/common/src/main/java/dev/cel/common/CelSource.java @@ -279,6 +279,10 @@ public Builder clearMacroCall(long exprId) { return this; } + public ImmutableSet getExtensions() { + return extensions.build(); + } + /** * Adds one or more {@link Extension}s to the source information. Extensions implement set * semantics and deduped if same ones are provided. diff --git a/common/src/main/java/dev/cel/common/navigation/BUILD.bazel b/common/src/main/java/dev/cel/common/navigation/BUILD.bazel index 1c1cd181..903ae3f0 100644 --- a/common/src/main/java/dev/cel/common/navigation/BUILD.bazel +++ b/common/src/main/java/dev/cel/common/navigation/BUILD.bazel @@ -57,6 +57,7 @@ java_library( ":common", "//:auto_value", "//common/ast:mutable_ast", + "//common/types:type_providers", "@maven//:com_google_errorprone_error_prone_annotations", "@maven//:com_google_guava_guava", "@maven//:org_jspecify_jspecify", diff --git a/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java b/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java index f28ada79..bda18dc0 100644 --- a/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java +++ b/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java @@ -15,6 +15,8 @@ package dev.cel.common.navigation; import dev.cel.common.ast.CelMutableAst; +import dev.cel.common.types.CelType; +import java.util.Optional; /** * Decorates a {@link CelMutableAst} with navigational properties. This allows us to visit a node's @@ -44,4 +46,15 @@ public CelNavigableMutableExpr getRoot() { public CelMutableAst getAst() { return ast; } + + /** + * Returns the type of the expression node for a type-checked AST. This simply proxies down the + * call to {@link CelMutableAst#getType(long)}. + * + * @return Optional of {@link CelType} or {@link Optional#empty} if the type does not exist at the + * ID. + */ + public Optional getType(long exprId) { + return ast.getType(exprId); + } } diff --git a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java index 94a4f472..2ce5fddd 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java +++ b/optimizer/src/main/java/dev/cel/optimizer/AstMutator.java @@ -14,37 +14,39 @@ package dev.cel.optimizer; -import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static java.lang.Math.max; +import static java.util.stream.Collectors.toCollection; import com.google.auto.value.AutoValue; -import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.HashBasedTable; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Table; import com.google.errorprone.annotations.Immutable; import dev.cel.common.CelAbstractSyntaxTree; import dev.cel.common.CelSource; +import dev.cel.common.ast.CelConstant; import dev.cel.common.ast.CelExpr; -import dev.cel.common.ast.CelExpr.CelCall; -import dev.cel.common.ast.CelExpr.CelComprehension; -import dev.cel.common.ast.CelExpr.CelIdent; import dev.cel.common.ast.CelExpr.ExprKind.Kind; -import dev.cel.common.ast.CelExprFactory; import dev.cel.common.ast.CelExprIdGeneratorFactory; import dev.cel.common.ast.CelExprIdGeneratorFactory.ExprIdGenerator; -import dev.cel.common.ast.CelExprIdGeneratorFactory.MonotonicIdGenerator; import dev.cel.common.ast.CelExprIdGeneratorFactory.StableIdGenerator; -import dev.cel.common.navigation.CelNavigableAst; +import dev.cel.common.ast.CelMutableAst; +import dev.cel.common.ast.CelMutableExpr; +import dev.cel.common.ast.CelMutableExpr.CelMutableCall; +import dev.cel.common.ast.CelMutableExpr.CelMutableComprehension; +import dev.cel.common.ast.CelMutableExpr.CelMutableCreateList; +import dev.cel.common.ast.CelMutableExprConverter; import dev.cel.common.navigation.CelNavigableExpr; +import dev.cel.common.navigation.CelNavigableMutableAst; +import dev.cel.common.navigation.CelNavigableMutableExpr; import dev.cel.common.navigation.TraversalOrder; import dev.cel.common.types.CelType; -import java.util.Collection; +import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map.Entry; import java.util.NoSuchElementException; import java.util.Optional; @@ -53,12 +55,14 @@ /** AstMutator contains logic for mutating a {@link CelAbstractSyntaxTree}. */ @Immutable +@SuppressWarnings("InlineMeSuggester") // Deprecated methods are not used and will be removed. public final class AstMutator { private static final ExprIdGenerator NO_OP_ID_GENERATOR = id -> id; + private static final ExprIdGenerator UNSET_ID_GENERATOR = id -> 0; private final long iterationLimit; /** - * Returns a new instance of a AST Mutator with the iteration limit set. + * Returns a new instance of a AST mutator with the iteration limit set. * *

Mutation is performed by walking the existing AST until the expression node to replace is * found, then the new subtree is walked to complete the mutation. Visiting of each node @@ -76,86 +80,70 @@ private AstMutator(long iterationLimit) { this.iterationLimit = iterationLimit; } - /** Replaces all the expression IDs in the expression tree with 0. */ - public CelExpr clearExprIds(CelExpr celExpr) { - return renumberExprIds((unused) -> 0, celExpr.toBuilder()).build(); - } - /** - * Replaces a subtree in the given expression node. This operation is intended for AST - * optimization purposes. - * - *

This is a very dangerous operation. Callers should 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. - * - *

If the ability to unparse an expression containing a macro call must be retained, use {@link - * #replaceSubtree(CelAbstractSyntaxTree, CelExpr, long) instead.} + * TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. * - * @param celExpr Original expression node to rewrite. - * @param newExpr New CelExpr to replace the subtree with. - * @param exprIdToReplace Expression id of the subtree that is getting replaced. + * @deprecated Use MutableExpr based APIs instead. */ - public CelExpr replaceSubtree(CelExpr celExpr, CelExpr newExpr, long exprIdToReplace) { - MonotonicIdGenerator monotonicIdGenerator = - CelExprIdGeneratorFactory.newMonotonicIdGenerator(0); - return mutateExpr( - unused -> monotonicIdGenerator.nextExprId(), - celExpr.toBuilder(), - newExpr.toBuilder(), - exprIdToReplace) - .build(); + @Deprecated + public CelExpr clearExprIds(CelExpr expr) { + return CelMutableExprConverter.fromMutableExpr( + clearExprIds(CelMutableExprConverter.fromCelExpr(expr))); + } + + /** Replaces all the expression IDs in the expression tree with 0. */ + public CelMutableExpr clearExprIds(CelMutableExpr expr) { + return renumberExprIds(UNSET_ID_GENERATOR, expr); } /** - * Replaces a subtree in the given AST. This operation is intended for AST optimization purposes. + * TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. * - *

This is a very dangerous operation. Callers should 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. - * - *

This will scrub out the description, positions and line offsets from {@code CelSource}. If - * the source contains macro calls, its call IDs will be to be consistent with the renumbered IDs - * in the AST. - * - * @param ast Original ast to mutate. - * @param newExpr New CelExpr to replace the subtree with. - * @param exprIdToReplace Expression id of the subtree that is getting replaced. + * @deprecated Use MutableExpr based APIs instead. */ - public CelAbstractSyntaxTree replaceSubtree( - CelAbstractSyntaxTree ast, CelExpr newExpr, long exprIdToReplace) { - return replaceSubtreeWithNewAst( - ast, - CelAbstractSyntaxTree.newParsedAst( - newExpr, - // Copy the macro call information to the new AST such that macro call map can be - // normalized post-replacement. - CelSource.newBuilder().addAllMacroCalls(ast.getSource().getMacroCalls()).build()), - exprIdToReplace); + @Deprecated + public CelAbstractSyntaxTree wrapAstWithNewCelBlock( + String celBlockFunction, CelAbstractSyntaxTree ast, List subexpressions) { + return wrapAstWithNewCelBlock( + celBlockFunction, + CelMutableAst.fromCelAst(ast), + subexpressions.stream() + .map(CelMutableExprConverter::fromCelExpr) + .collect(toCollection(ArrayList::new))) + .toParsedAst(); } /** Wraps the given AST and its subexpressions with a new cel.@block call. */ - public CelAbstractSyntaxTree wrapAstWithNewCelBlock( - String celBlockFunction, CelAbstractSyntaxTree ast, Collection subexpressions) { + public CelMutableAst wrapAstWithNewCelBlock( + String celBlockFunction, CelMutableAst ast, List subexpressions) { long maxId = getMaxId(ast); - CelExpr blockExpr = - CelExpr.newBuilder() - .setId(++maxId) - .setCall( - CelCall.newBuilder() - .setFunction(celBlockFunction) - .addArgs( - CelExpr.ofCreateListExpr( - ++maxId, ImmutableList.copyOf(subexpressions), ImmutableList.of()), - ast.getExpr()) - .build()) - .build(); - - return CelAbstractSyntaxTree.newParsedAst(blockExpr, ast.getSource()); + CelMutableExpr blockExpr = + CelMutableExpr.ofCall( + ++maxId, + CelMutableCall.create( + celBlockFunction, + CelMutableExpr.ofCreateList(++maxId, CelMutableCreateList.create(subexpressions)), + ast.expr())); + + return CelMutableAst.of(blockExpr, ast.source()); + } + + /** + * TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. + * + * @deprecated Use MutableExpr based APIs instead. + */ + @Deprecated + public CelAbstractSyntaxTree replaceSubtreeWithNewBindMacro( + CelAbstractSyntaxTree ast, String varName, CelExpr varInit, CelExpr resultExpr, long id) { + return replaceSubtreeWithNewBindMacro( + CelMutableAst.fromCelAst(ast), + varName, + CelMutableExprConverter.fromCelExpr(varInit), + CelMutableExprConverter.fromCelExpr(resultExpr), + id, + true) + .toParsedAst(); } /** @@ -164,49 +152,62 @@ public CelAbstractSyntaxTree wrapAstWithNewCelBlock( * *

The bind call takes the format of: {@code cel.bind(varInit, varName, resultExpr)} * - * @param ast Original ast to mutate. + * @param ast Original AST to mutate. * @param varName New variable name for the bind macro call. * @param varInit Initialization expression to bind to the local variable. * @param resultExpr Result expression * @param exprIdToReplace Expression ID of the subtree that is getting replaced. + * @param populateMacroSource If true, populates the cel.bind macro source in the AST. */ - public CelAbstractSyntaxTree replaceSubtreeWithNewBindMacro( - CelAbstractSyntaxTree ast, + public CelMutableAst replaceSubtreeWithNewBindMacro( + CelMutableAst ast, String varName, - CelExpr varInit, - CelExpr resultExpr, - long exprIdToReplace) { + CelMutableExpr varInit, + CelMutableExpr resultExpr, + long exprIdToReplace, + boolean populateMacroSource) { + // Copy the incoming expressions to prevent modifying the root long maxId = max(getMaxId(varInit), getMaxId(ast)); StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(maxId); - BindMacro bindMacro = newBindMacro(varName, varInit, resultExpr, stableIdGenerator); - // In situations where the existing AST already contains a macro call (ex: nested cel.binds), - // its macro source must be normalized to make it consistent with the newly generated bind - // macro. - CelSource celSource = - normalizeMacroSource( - ast.getSource(), - -1, // Do not replace any of the subexpr in the macro map. - bindMacro.bindMacro().toBuilder(), - stableIdGenerator::renumberId); - celSource = - celSource.toBuilder() - .addMacroCalls(bindMacro.bindExpr().id(), bindMacro.bindMacro()) - .build(); + CelMutableExpr newBindMacroExpr = + newBindMacroExpr( + varName, varInit, CelMutableExpr.newInstance(resultExpr), stableIdGenerator); + CelSource.Builder celSource = CelSource.newBuilder(); + if (populateMacroSource) { + CelMutableExpr newBindMacroSourceExpr = + newBindMacroSourceExpr(newBindMacroExpr, varName, stableIdGenerator); + // In situations where the existing AST already contains a macro call (ex: nested cel.binds), + // its macro source must be normalized to make it consistent with the newly generated bind + // macro. + celSource = + normalizeMacroSource( + ast.source(), + -1, // Do not replace any of the subexpr in the macro map. + newBindMacroSourceExpr, + stableIdGenerator::renumberId) + .addMacroCalls( + newBindMacroExpr.id(), + CelMutableExprConverter.fromMutableExpr(newBindMacroSourceExpr)); + } + + CelMutableAst newBindAst = CelMutableAst.of(newBindMacroExpr, celSource); - return replaceSubtreeWithNewAst( - ast, CelAbstractSyntaxTree.newParsedAst(bindMacro.bindExpr(), celSource), exprIdToReplace); + return replaceSubtree(ast, newBindAst, exprIdToReplace); } - /** Renumbers all the expr IDs in the given AST in a consecutive manner starting from 1. */ public CelAbstractSyntaxTree renumberIdsConsecutively(CelAbstractSyntaxTree ast) { + return renumberIdsConsecutively(CelMutableAst.fromCelAst(ast)).toParsedAst(); + } + + /** Renumbers all the expr IDs in the given AST in a consecutive manner starting from 1. */ + public CelMutableAst renumberIdsConsecutively(CelMutableAst mutableAst) { StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(0); - CelExpr.Builder root = - renumberExprIds(stableIdGenerator::renumberId, ast.getExpr().toBuilder()); - CelSource newSource = + CelMutableExpr mutableExpr = renumberExprIds(stableIdGenerator::renumberId, mutableAst.expr()); + CelSource.Builder newSource = normalizeMacroSource( - ast.getSource(), Integer.MIN_VALUE, root, stableIdGenerator::renumberId); + mutableAst.source(), Integer.MIN_VALUE, mutableExpr, stableIdGenerator::renumberId); - return CelAbstractSyntaxTree.newParsedAst(root.build(), newSource); + return CelMutableAst.of(mutableExpr, newSource); } /** @@ -235,23 +236,24 @@ public CelAbstractSyntaxTree renumberIdsConsecutively(CelAbstractSyntaxTree ast) * && [true].exists(@c1:0, @c1:0))} * * - * @param ast AST to mutate + * @param ast AST containing type-checked references * @param newIterVarPrefix Prefix to use for new iteration variable identifier name. For example, * providing @c will produce @c0:0, @c0:1, @c1:0, @c2:0... as new names. * @param newResultPrefix Prefix to use for new comprehensin result identifier names. */ public MangledComprehensionAst mangleComprehensionIdentifierNames( CelAbstractSyntaxTree ast, String newIterVarPrefix, String newResultPrefix) { - CelNavigableAst newNavigableAst = CelNavigableAst.fromAst(ast); - Predicate comprehensionIdentifierPredicate = x -> true; + CelNavigableMutableAst navigableMutableAst = + CelNavigableMutableAst.fromAst(CelMutableAst.fromCelAst(ast)); + Predicate comprehensionIdentifierPredicate = x -> true; comprehensionIdentifierPredicate = comprehensionIdentifierPredicate .and(node -> node.getKind().equals(Kind.COMPREHENSION)) .and(node -> !node.expr().comprehension().iterVar().startsWith(newIterVarPrefix)) .and(node -> !node.expr().comprehension().accuVar().startsWith(newResultPrefix)); - LinkedHashMap comprehensionsToMangle = - newNavigableAst + LinkedHashMap comprehensionsToMangle = + navigableMutableAst .getRoot() // This is important - mangling needs to happen bottom-up to avoid stepping over // shadowed variables that are not part of the comprehension being mangled. @@ -260,11 +262,10 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames( .filter( node -> { // Ensure the iter_var or the comprehension result is actually referenced in the - // loop_step. If it's not, we - // can skip mangling. + // loop_step. If it's not, we can skip mangling. String iterVar = node.expr().comprehension().iterVar(); String result = node.expr().comprehension().result().ident().name(); - return CelNavigableExpr.fromExpr(node.expr().comprehension().loopStep()) + return CelNavigableMutableExpr.fromExpr(node.expr().comprehension().loopStep()) .allNodes() .filter(subNode -> subNode.getKind().equals(Kind.IDENT)) .map(subNode -> subNode.expr().ident()) @@ -275,36 +276,46 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames( Collectors.toMap( k -> k, v -> { - CelComprehension comprehension = v.expr().comprehension(); + CelMutableComprehension comprehension = v.expr().comprehension(); String iterVar = comprehension.iterVar(); - // Identifiers to mangle could be the iteration variable, comprehension result - // or both, but at least one has to exist. - // As an example, [1,2].map(i, 3) would produce an optional.empty because `i` - // is not actually used. + // Identifiers to mangle could be the iteration variable, comprehension + // result or both, but at least one has to exist. + // As an example, [1,2].map(i, 3) would result in optional.empty for iteration + // variable because `i` is not actually used. Optional iterVarId = - CelNavigableExpr.fromExpr(comprehension.loopStep()) + CelNavigableMutableExpr.fromExpr(comprehension.loopStep()) .allNodes() .filter( loopStepNode -> - loopStepNode.expr().identOrDefault().name().equals(iterVar)) - .map(CelNavigableExpr::id) + loopStepNode.getKind().equals(Kind.IDENT) + && loopStepNode.expr().ident().name().equals(iterVar)) + .map(CelNavigableMutableExpr::id) .findAny(); Optional iterVarType = iterVarId.map( id -> - ast.getType(id) + navigableMutableAst + .getType(id) .orElseThrow( () -> new NoSuchElementException( - "Checked type not present for iteration variable:" - + " " + "Checked type not present for iteration" + + " variable: " + iterVarId))); - Optional resultType = ast.getType(comprehension.result().id()); + CelType resultType = + navigableMutableAst + .getType(comprehension.result().id()) + .orElseThrow( + () -> + new IllegalStateException( + "Result type was not present for the comprehension ID: " + + comprehension.result().id())); return MangledComprehensionType.of(iterVarType, resultType); }, (x, y) -> { - throw new IllegalStateException("Unexpected CelNavigableExpr collision"); + throw new IllegalStateException( + "Unexpected CelNavigableMutableExpr collision"); }, LinkedHashMap::new)); int iterCount = 0; @@ -315,21 +326,15 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames( // Intermediary table used for the purposes of generating a unique mangled variable name. Table comprehensionLevelToType = HashBasedTable.create(); - for (Entry comprehensionEntry : + CelMutableExpr mutatedComprehensionExpr = navigableMutableAst.getAst().expr(); + CelSource.Builder newSource = navigableMutableAst.getAst().source(); + for (Entry comprehensionEntry : comprehensionsToMangle.entrySet()) { iterCount++; - // Refetch the comprehension node as mutating the AST could have renumbered its IDs. - CelNavigableExpr comprehensionNode = - newNavigableAst - .getRoot() - .allNodes(TraversalOrder.POST_ORDER) - .filter(comprehensionIdentifierPredicate) - .findAny() - .orElseThrow( - () -> new NoSuchElementException("Failed to refetch mutated comprehension")); + CelNavigableMutableExpr comprehensionNode = comprehensionEntry.getKey(); MangledComprehensionType comprehensionEntryType = comprehensionEntry.getValue(); - CelExpr.Builder comprehensionExpr = comprehensionNode.expr().toBuilder(); + CelMutableExpr comprehensionExpr = comprehensionNode.expr(); int comprehensionNestingLevel = countComprehensionNestingLevel(comprehensionNode); MangledComprehensionName mangledComprehensionName; if (comprehensionLevelToType.contains(comprehensionNestingLevel, comprehensionEntryType)) { @@ -352,25 +357,21 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames( String iterVar = comprehensionExpr.comprehension().iterVar(); String accuVar = comprehensionExpr.comprehension().accuVar(); - CelExpr.Builder mutatedComprehensionExpr = + mutatedComprehensionExpr = mangleIdentsInComprehensionExpr( - newNavigableAst.getAst().getExpr().toBuilder(), + mutatedComprehensionExpr, comprehensionExpr, iterVar, accuVar, mangledComprehensionName); // Repeat the mangling process for the macro source. - CelSource newSource = + newSource = mangleIdentsInMacroSource( - newNavigableAst.getAst(), + newSource, mutatedComprehensionExpr, iterVar, mangledComprehensionName, comprehensionExpr.id()); - - newNavigableAst = - CelNavigableAst.fromAst( - CelAbstractSyntaxTree.newParsedAst(mutatedComprehensionExpr.build(), newSource)); } if (iterCount >= iterationLimit) { @@ -381,100 +382,200 @@ public MangledComprehensionAst mangleComprehensionIdentifierNames( } return MangledComprehensionAst.of( - newNavigableAst.getAst(), ImmutableMap.copyOf(mangledIdentNamesToType)); + CelMutableAst.of(mutatedComprehensionExpr, newSource), + ImmutableMap.copyOf(mangledIdentNamesToType)); + } + + /** + * TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. + * + * @deprecated Use MutableExpr based APIs instead. + */ + @Deprecated + public CelExpr replaceSubtree(CelExpr celExpr, CelExpr newExpr, long id) { + return replaceSubtree( + CelMutableExprConverter.fromCelExpr(celExpr), + CelMutableExprConverter.fromCelExpr(newExpr), + id) + .toParsedAst() + .getExpr(); + } + + /** + * TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. + * + * @deprecated Use MutableExpr based APIs instead. + */ + @Deprecated + public CelAbstractSyntaxTree replaceSubtree( + CelAbstractSyntaxTree root, CelExpr newExpr, long id) { + return replaceSubtree( + CelMutableAst.fromCelAst(root), CelMutableExprConverter.fromCelExpr(newExpr), id) + .toParsedAst(); + } + + /** + * Replaces a subtree in the given expression node. 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. + * + *

If the ability to unparse an expression containing a macro call must be retained, use {@link + * #replaceSubtree(CelMutableAst, CelMutableAst, long) instead.} + * + * @param root Original expression node to rewrite. + * @param newExpr New CelExpr to replace the subtree with. + * @param exprIdToReplace Expression id of the subtree that is getting replaced. + */ + public CelMutableAst replaceSubtree( + CelMutableExpr root, CelMutableExpr newExpr, long exprIdToReplace) { + return replaceSubtree(CelMutableAst.of(root, CelSource.newBuilder()), newExpr, exprIdToReplace); + } + + /** + * Replaces a subtree in the given AST. 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. + * + *

This will scrub out the description, positions and line offsets from {@code CelSource}. If + * the source contains macro calls, its call IDs will be to be consistent with the renumbered IDs + * in the AST. + * + * @param ast Original ast to mutate. + * @param newExpr New CelExpr to replace the subtree with. + * @param exprIdToReplace Expression id of the subtree that is getting replaced. + */ + public CelMutableAst replaceSubtree( + CelMutableAst ast, CelMutableExpr newExpr, long exprIdToReplace) { + return replaceSubtree( + ast, + CelMutableAst.of( + newExpr, + // Copy the macro call information to the new AST such that macro call map can be + // normalized post-replacement. + CelSource.newBuilder().addAllMacroCalls(ast.source().getMacroCalls())), + exprIdToReplace); } /** - * Mutates the given AST by replacing a subtree at a given index. + * Replaces a subtree in the given AST. This operation is intended for AST optimization purposes. * - * @param ast Existing AST being mutated - * @param newAst New subtree to perform the replacement with. If the subtree has a macro map - * populated, its macro source is merged with the existing AST's after normalization. - * @param exprIdToReplace The expr ID in the existing AST to replace the subtree at. + *

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. + * + *

This will scrub out the description, positions and line offsets from {@code CelSource}. If + * the source contains macro calls, its call IDs will be to be consistent with the renumbered IDs + * in the AST. + * + * @param ast Original ast to mutate. + * @param newAst New AST to replace the subtree with. + * @param exprIdToReplace Expression id of the subtree that is getting replaced. */ - @VisibleForTesting - CelAbstractSyntaxTree replaceSubtreeWithNewAst( - CelAbstractSyntaxTree ast, CelAbstractSyntaxTree newAst, long exprIdToReplace) { + public CelMutableAst replaceSubtree( + CelMutableAst ast, CelMutableAst newAst, long exprIdToReplace) { + return replaceSubtree( + CelNavigableMutableAst.fromAst(ast), + CelNavigableMutableAst.fromAst(newAst), + exprIdToReplace); + } + + /** + * Replaces a subtree in the given AST. 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. + * + *

This will scrub out the description, positions and line offsets from {@code CelSource}. If + * the source contains macro calls, its call IDs will be to be consistent with the renumbered IDs + * in the AST. + * + * @param navAst Original navigable ast to mutate. + * @param navNewAst New navigable AST to replace the subtree with. + * @param exprIdToReplace Expression id of the subtree that is getting replaced. + */ + public CelMutableAst replaceSubtree( + CelNavigableMutableAst navAst, CelNavigableMutableAst navNewAst, long exprIdToReplace) { // Stabilize the incoming AST by renumbering all of its expression IDs. - long maxId = max(getMaxId(ast), getMaxId(newAst)); + long maxId = max(getMaxId(navAst), getMaxId(navNewAst)); + CelMutableAst ast = navAst.getAst(); + CelMutableAst newAst = navNewAst.getAst(); newAst = stabilizeAst(newAst, maxId); + long stablizedNewExprRootId = newAst.expr().id(); // Mutate the AST root with the new subtree. All the existing expr IDs are renumbered in the // process, but its original IDs are memoized so that we can normalize the expr IDs // in the macro source map. StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(getMaxId(newAst)); - CelExpr.Builder mutatedRoot = - mutateExpr( - stableIdGenerator::renumberId, - ast.getExpr().toBuilder(), - newAst.getExpr().toBuilder(), - exprIdToReplace); - - CelSource newAstSource = ast.getSource(); - if (!newAst.getSource().getMacroCalls().isEmpty()) { - // The root is mutated, but the expr IDs in the macro map needs to be normalized. - // In situations where an AST with a new macro map is being inserted (ex: new bind call), - // the new subtree's expr ID is not memoized in the stable ID generator because the ID never - // existed in the main AST. - // In this case, we forcibly memoize the new subtree ID with a newly generated ID so - // that the macro map IDs can be normalized properly. + + CelMutableExpr mutatedRoot = + mutateExpr(stableIdGenerator::renumberId, ast.expr(), newAst.expr(), exprIdToReplace); + CelSource.Builder newAstSource = CelSource.newBuilder(); + if (!ast.source().getMacroCalls().isEmpty()) { + newAstSource = combine(newAstSource, ast.source()); + } + + if (!newAst.source().getMacroCalls().isEmpty()) { stableIdGenerator.memoize( - newAst.getExpr().id(), stableIdGenerator.renumberId(exprIdToReplace)); - newAstSource = combine(newAstSource, newAst.getSource()); + stablizedNewExprRootId, stableIdGenerator.renumberId(exprIdToReplace)); + newAstSource = combine(newAstSource, newAst.source()); } newAstSource = normalizeMacroSource( newAstSource, exprIdToReplace, mutatedRoot, stableIdGenerator::renumberId); - - return CelAbstractSyntaxTree.newParsedAst(mutatedRoot.build(), newAstSource); + return CelMutableAst.of(mutatedRoot, newAstSource); } - private CelExpr.Builder mangleIdentsInComprehensionExpr( - CelExpr.Builder root, - CelExpr.Builder comprehensionExpr, + private CelMutableExpr mangleIdentsInComprehensionExpr( + CelMutableExpr root, + CelMutableExpr comprehensionExpr, String originalIterVar, String originalAccuVar, MangledComprehensionName mangledComprehensionName) { - CelExpr.Builder modifiedLoopStep = - replaceIdentName( - comprehensionExpr.comprehension().loopStep().toBuilder(), - originalIterVar, - mangledComprehensionName.iterVarName()); - comprehensionExpr.setComprehension( - comprehensionExpr.comprehension().toBuilder() - .setLoopStep(modifiedLoopStep.build()) - .build()); - comprehensionExpr = - replaceIdentName(comprehensionExpr, originalAccuVar, mangledComprehensionName.resultName()); - - CelComprehension.Builder newComprehension = - comprehensionExpr.comprehension().toBuilder() - .setIterVar(mangledComprehensionName.iterVarName()); + CelMutableComprehension comprehension = comprehensionExpr.comprehension(); + replaceIdentName( + comprehension.loopStep(), originalIterVar, mangledComprehensionName.iterVarName()); + replaceIdentName(comprehensionExpr, originalAccuVar, mangledComprehensionName.resultName()); + + comprehension.setIterVar(mangledComprehensionName.iterVarName()); // Most standard macros set accu_var as __result__, but not all (ex: cel.bind). - if (newComprehension.accuVar().equals(originalAccuVar)) { - newComprehension.setAccuVar(mangledComprehensionName.resultName()); + if (comprehension.accuVar().equals(originalAccuVar)) { + comprehension.setAccuVar(mangledComprehensionName.resultName()); } - return mutateExpr( - NO_OP_ID_GENERATOR, - root, - comprehensionExpr.setComprehension(newComprehension.build()), - comprehensionExpr.id()); + return mutateExpr(NO_OP_ID_GENERATOR, root, comprehensionExpr, comprehensionExpr.id()); } - private CelExpr.Builder replaceIdentName( - CelExpr.Builder comprehensionExpr, String originalIdentName, String newIdentName) { + private void replaceIdentName( + CelMutableExpr comprehensionExpr, String originalIdentName, String newIdentName) { int iterCount; for (iterCount = 0; iterCount < iterationLimit; iterCount++) { - Optional identToMangle = - CelNavigableExpr.fromExpr(comprehensionExpr.build()) + CelMutableExpr identToMangle = + CelNavigableMutableExpr.fromExpr(comprehensionExpr) .descendants() - .map(CelNavigableExpr::expr) - .filter(node -> node.identOrDefault().name().equals(originalIdentName)) - .findAny(); - if (!identToMangle.isPresent()) { + .map(CelNavigableMutableExpr::expr) + .filter( + node -> + node.getKind().equals(Kind.IDENT) + && node.ident().name().equals(originalIdentName)) + .findAny() + .orElse(null); + if (identToMangle == null) { break; } @@ -482,31 +583,29 @@ private CelExpr.Builder replaceIdentName( mutateExpr( NO_OP_ID_GENERATOR, comprehensionExpr, - CelExpr.newBuilder().setIdent(CelIdent.newBuilder().setName(newIdentName).build()), - identToMangle.get().id()); + CelMutableExpr.ofIdent(newIdentName), + identToMangle.id()); } if (iterCount >= iterationLimit) { throw new IllegalStateException("Max iteration count reached."); } - - return comprehensionExpr; } - private CelSource mangleIdentsInMacroSource( - CelAbstractSyntaxTree ast, - CelExpr.Builder mutatedComprehensionExpr, + private CelSource.Builder mangleIdentsInMacroSource( + CelSource.Builder sourceBuilder, + CelMutableExpr mutatedComprehensionExpr, String originalIterVar, MangledComprehensionName mangledComprehensionName, long originalComprehensionId) { - if (!ast.getSource().getMacroCalls().containsKey(originalComprehensionId)) { - return ast.getSource(); + if (!sourceBuilder.getMacroCalls().containsKey(originalComprehensionId)) { + return sourceBuilder; } // First, normalize the macro source. // ex: [x].exists(x, [x].exists(x, x == 1)) -> [x].exists(x, [@c1].exists(x, @c0 == 1)). CelSource.Builder newSource = - normalizeMacroSource(ast.getSource(), -1, mutatedComprehensionExpr, (id) -> id).toBuilder(); + normalizeMacroSource(sourceBuilder, -1, mutatedComprehensionExpr, (id) -> id); // Note that in the above example, the iteration variable is not replaced after normalization. // This is because populating a macro call map upon parse generates a new unique identifier @@ -515,62 +614,69 @@ private CelSource mangleIdentsInMacroSource( // variable actually exists in the main AST thus, this step isn't needed. // ex: [1].map(x, [2].filter(y, x == y). Here, the variable declaration `x` exists in the AST // but not `y`. - CelExpr.Builder macroExpr = newSource.getMacroCalls().get(originalComprehensionId).toBuilder(); + CelMutableExpr macroExpr = + CelMutableExprConverter.fromCelExpr(newSource.getMacroCalls().get(originalComprehensionId)); // By convention, the iteration variable is always the first argument of the // macro call expression. - CelExpr identToMangle = macroExpr.call().args().get(0); - if (identToMangle.identOrDefault().name().equals(originalIterVar)) { + CelMutableExpr identToMangle = macroExpr.call().args().get(0); + if (identToMangle.ident().name().equals(originalIterVar)) { + // if (identToMangle.identOrDefault().name().equals(originalIterVar)) { macroExpr = mutateExpr( NO_OP_ID_GENERATOR, macroExpr, - CelExpr.newBuilder() - .setIdent( - CelIdent.newBuilder() - .setName(mangledComprehensionName.iterVarName()) - .build()), + CelMutableExpr.ofIdent(mangledComprehensionName.iterVarName()), identToMangle.id()); } - newSource.addMacroCalls(originalComprehensionId, macroExpr.build()); - return newSource.build(); + newSource.addMacroCalls( + originalComprehensionId, CelMutableExprConverter.fromMutableExpr(macroExpr)); + + return newSource; } - private BindMacro newBindMacro( - String varName, CelExpr varInit, CelExpr resultExpr, StableIdGenerator stableIdGenerator) { + private CelMutableExpr newBindMacroExpr( + String varName, + CelMutableExpr varInit, + CelMutableExpr resultExpr, + StableIdGenerator stableIdGenerator) { // Renumber incoming expression IDs in the init and result expression to avoid collision with // the main AST. Existing IDs are memoized for a macro source sanitization pass at the end // (e.g: inserting a bind macro to an existing macro expr) - varInit = renumberExprIds(stableIdGenerator::nextExprId, varInit.toBuilder()).build(); - resultExpr = renumberExprIds(stableIdGenerator::nextExprId, resultExpr.toBuilder()).build(); - CelExprFactory exprFactory = - CelExprFactory.newInstance((unused) -> stableIdGenerator.nextExprId()); - CelExpr bindMacroExpr = - exprFactory.fold( + varInit = renumberExprIds(stableIdGenerator::nextExprId, varInit); + resultExpr = renumberExprIds(stableIdGenerator::nextExprId, resultExpr); + + long iterRangeId = stableIdGenerator.nextExprId(); + long loopConditionId = stableIdGenerator.nextExprId(); + long loopStepId = stableIdGenerator.nextExprId(); + long comprehensionId = stableIdGenerator.nextExprId(); + + return CelMutableExpr.ofComprehension( + comprehensionId, + CelMutableComprehension.create( "#unused", - exprFactory.newList(), + CelMutableExpr.ofCreateList(iterRangeId, CelMutableCreateList.create()), varName, varInit, - exprFactory.newBoolLiteral(false), - exprFactory.newIdentifier(varName), - resultExpr); - - CelExpr bindMacroCallExpr = - exprFactory - .newReceiverCall( - "bind", - CelExpr.ofIdentExpr(stableIdGenerator.nextExprId(), "cel"), - CelExpr.ofIdentExpr(stableIdGenerator.nextExprId(), varName), - bindMacroExpr.comprehension().accuInit(), - bindMacroExpr.comprehension().result()) - .toBuilder() - .setId(0) - .build(); - - return BindMacro.of(bindMacroExpr, bindMacroCallExpr); - } - - private static CelSource combine(CelSource celSource1, CelSource celSource2) { + CelMutableExpr.ofConstant(loopConditionId, CelConstant.ofValue(false)), + CelMutableExpr.ofIdent(loopStepId, varName), + resultExpr)); + } + + private CelMutableExpr newBindMacroSourceExpr( + CelMutableExpr bindMacroExpr, String varName, StableIdGenerator stableIdGenerator) { + return CelMutableExpr.ofCall( + 0, // Required sentinel value for macro call + CelMutableCall.create( + CelMutableExpr.ofIdent(stableIdGenerator.nextExprId(), "cel"), + "bind", + CelMutableExpr.ofIdent(stableIdGenerator.nextExprId(), varName), + bindMacroExpr.comprehension().accuInit(), + bindMacroExpr.comprehension().result())); + } + + private static CelSource.Builder combine( + CelSource.Builder celSource1, CelSource.Builder celSource2) { ImmutableMap.Builder macroMap = ImmutableMap.builder(); macroMap.putAll(celSource1.getMacroCalls()); macroMap.putAll(celSource2.getMacroCalls()); @@ -578,8 +684,7 @@ private static CelSource combine(CelSource celSource1, CelSource celSource2) { return CelSource.newBuilder() .addAllExtensions(celSource1.getExtensions()) .addAllExtensions(celSource2.getExtensions()) - .addAllMacroCalls(macroMap.buildOrThrow()) - .build(); + .addAllMacroCalls(macroMap.buildOrThrow()); } /** @@ -587,52 +692,51 @@ private static CelSource combine(CelSource celSource1, CelSource celSource2) { * (monotonically increased) from the starting seed ID. If the AST contains any macro calls, its * IDs are also normalized. */ - private CelAbstractSyntaxTree stabilizeAst(CelAbstractSyntaxTree ast, long seedExprId) { + private CelMutableAst stabilizeAst(CelMutableAst mutableAst, long seedExprId) { + CelMutableExpr mutableExpr = mutableAst.expr(); + CelSource.Builder source = mutableAst.source(); StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(seedExprId); - CelExpr.Builder newExprBuilder = - renumberExprIds(stableIdGenerator::nextExprId, ast.getExpr().toBuilder()); - - if (ast.getSource().getMacroCalls().isEmpty()) { - return CelAbstractSyntaxTree.newParsedAst(newExprBuilder.build(), ast.getSource()); - } + CelMutableExpr mutatedExpr = renumberExprIds(stableIdGenerator::nextExprId, mutableExpr); CelSource.Builder sourceBuilder = - CelSource.newBuilder().addAllExtensions(ast.getSource().getExtensions()); + CelSource.newBuilder().addAllExtensions(source.getExtensions()); // Update the macro call IDs and their call IDs - for (Entry macroCall : ast.getSource().getMacroCalls().entrySet()) { + for (Entry macroCall : source.getMacroCalls().entrySet()) { long macroId = macroCall.getKey(); long newCallId = stableIdGenerator.renumberId(macroId); + CelMutableExpr existingMacroCallExpr = + CelMutableExprConverter.fromCelExpr(macroCall.getValue()); - CelExpr.Builder newCall = - renumberExprIds(stableIdGenerator::renumberId, macroCall.getValue().toBuilder()); + CelMutableExpr newCall = + renumberExprIds(stableIdGenerator::renumberId, existingMacroCallExpr); - sourceBuilder.addMacroCalls(newCallId, newCall.build()); + sourceBuilder.addMacroCalls(newCallId, CelMutableExprConverter.fromMutableExpr(newCall)); } - return CelAbstractSyntaxTree.newParsedAst(newExprBuilder.build(), sourceBuilder.build()); + return CelMutableAst.of(mutatedExpr, sourceBuilder); } - private CelSource normalizeMacroSource( - CelSource celSource, + private CelSource.Builder normalizeMacroSource( + CelSource.Builder celSource, long exprIdToReplace, - CelExpr.Builder mutatedRoot, + CelMutableExpr mutatedRoot, ExprIdGenerator idGenerator) { // Remove the macro metadata that no longer exists in the AST due to being replaced. - celSource = celSource.toBuilder().clearMacroCall(exprIdToReplace).build(); + celSource.clearMacroCall(exprIdToReplace); CelSource.Builder sourceBuilder = CelSource.newBuilder().addAllExtensions(celSource.getExtensions()); if (celSource.getMacroCalls().isEmpty()) { - return sourceBuilder.build(); + return sourceBuilder; } - ImmutableMap allExprs = - CelNavigableExpr.fromExpr(mutatedRoot.build()) + ImmutableMap allExprs = + CelNavigableMutableExpr.fromExpr(mutatedRoot) .allNodes() - .map(CelNavigableExpr::expr) + .map(CelNavigableMutableExpr::expr) .collect( toImmutableMap( - CelExpr::id, + CelMutableExpr::id, expr -> expr, (expr1, expr2) -> { // Comprehensions can reuse same expression (result). We just need to ensure @@ -653,23 +757,26 @@ private CelSource normalizeMacroSource( continue; } - CelExpr.Builder newMacroCallExpr = - renumberExprIds(idGenerator, existingMacroCall.getValue().toBuilder()); + CelMutableExpr existingMacroCallExpr = + CelMutableExprConverter.fromCelExpr(existingMacroCall.getValue()); + CelMutableExpr newMacroCallExpr = renumberExprIds(idGenerator, existingMacroCallExpr); - CelNavigableExpr callNav = CelNavigableExpr.fromExpr(newMacroCallExpr.build()); - ImmutableList callDescendants = - callNav.descendants().map(CelNavigableExpr::expr).collect(toImmutableList()); + CelNavigableMutableExpr callNav = CelNavigableMutableExpr.fromExpr(newMacroCallExpr); + ArrayList callDescendants = + callNav + .descendants() + .map(CelNavigableMutableExpr::expr) + .collect(toCollection(ArrayList::new)); - for (CelExpr callChild : callDescendants) { + for (CelMutableExpr callChild : callDescendants) { if (!allExprs.containsKey(callChild.id())) { continue; } - CelExpr mutatedExpr = allExprs.get(callChild.id()); + CelMutableExpr mutatedExpr = allExprs.get(callChild.id()); if (!callChild.equals(mutatedExpr)) { newMacroCallExpr = - mutateExpr( - NO_OP_ID_GENERATOR, newMacroCallExpr, mutatedExpr.toBuilder(), callChild.id()); + mutateExpr(NO_OP_ID_GENERATOR, newMacroCallExpr, mutatedExpr, callChild.id()); } } @@ -677,55 +784,55 @@ private CelSource normalizeMacroSource( long replacedId = idGenerator.generate(exprIdToReplace); boolean isListExprBeingReplaced = allExprs.containsKey(replacedId) - && allExprs.get(replacedId).exprKind().getKind().equals(Kind.CREATE_LIST); + && allExprs.get(replacedId).getKind().equals(Kind.CREATE_LIST); if (isListExprBeingReplaced) { unwrapListArgumentsInMacroCallExpr( allExprs.get(callId).comprehension(), newMacroCallExpr); } } - sourceBuilder.addMacroCalls(callId, newMacroCallExpr.build()); + sourceBuilder.addMacroCalls( + callId, CelMutableExprConverter.fromMutableExpr(newMacroCallExpr)); } // Replace comprehension nodes with a NOT_SET reference to reduce AST size. for (Entry macroCall : sourceBuilder.getMacroCalls().entrySet()) { - CelExpr macroCallExpr = macroCall.getValue(); - CelNavigableExpr.fromExpr(macroCallExpr) + CelMutableExpr macroCallExpr = CelMutableExprConverter.fromCelExpr(macroCall.getValue()); + CelNavigableMutableExpr.fromExpr(macroCallExpr) .allNodes() .filter(node -> node.getKind().equals(Kind.COMPREHENSION)) - .map(CelNavigableExpr::expr) + .map(CelNavigableMutableExpr::expr) .forEach( node -> { - CelExpr.Builder mutatedNode = + CelMutableExpr mutatedNode = mutateExpr( - (id) -> id, - macroCallExpr.toBuilder(), - CelExpr.ofNotSet(node.id()).toBuilder(), + NO_OP_ID_GENERATOR, + macroCallExpr, + CelMutableExpr.ofNotSet(node.id()), node.id()); - macroCall.setValue(mutatedNode.build()); + macroCall.setValue(CelMutableExprConverter.fromMutableExpr(mutatedNode)); }); // Prune any NOT_SET (comprehension) nodes that no longer exist in the main AST // This can occur from pulling out a nested comprehension into a separate cel.block index - CelNavigableExpr.fromExpr(macroCallExpr) + CelNavigableMutableExpr.fromExpr(macroCallExpr) .allNodes() .filter(node -> node.getKind().equals(Kind.NOT_SET)) - .map(CelNavigableExpr::id) + .map(CelNavigableMutableExpr::id) .filter(id -> !allExprs.containsKey(id)) .forEach( id -> { - ImmutableList newCallArgs = + ArrayList newCallArgs = macroCallExpr.call().args().stream() .filter(node -> node.id() != id) - .collect(toImmutableList()); - CelCall.Builder call = - macroCallExpr.call().toBuilder().clearArgs().addArgs(newCallArgs); - - macroCall.setValue(macroCallExpr.toBuilder().setCall(call.build()).build()); + .collect(toCollection(ArrayList::new)); + CelMutableCall call = macroCallExpr.call(); + call.setArgs(newCallArgs); + macroCall.setValue(CelMutableExprConverter.fromMutableExpr(macroCallExpr)); }); } - return sourceBuilder.build(); + return sourceBuilder; } /** @@ -741,69 +848,83 @@ private CelSource normalizeMacroSource( * arguments unwrapped. */ private static void unwrapListArgumentsInMacroCallExpr( - CelComprehension comprehension, CelExpr.Builder newMacroCallExpr) { - CelExpr accuInit = comprehension.accuInit(); - if (!accuInit.exprKind().getKind().equals(Kind.CREATE_LIST) + CelMutableComprehension comprehension, CelMutableExpr newMacroCallExpr) { + CelMutableExpr accuInit = comprehension.accuInit(); + if (!accuInit.getKind().equals(Kind.CREATE_LIST) || !accuInit.createList().elements().isEmpty()) { // Does not contain an extraneous list. return; } - CelExpr loopStepExpr = comprehension.loopStep(); - ImmutableList args = loopStepExpr.call().args(); - if (args.size() != 2) { + CelMutableExpr loopStepExpr = comprehension.loopStep(); + List loopStepArgs = loopStepExpr.call().args(); + if (loopStepArgs.size() != 2) { throw new IllegalArgumentException( String.format( "Expected exactly 2 arguments but got %d instead on expr id: %d", - args.size(), loopStepExpr.id())); + loopStepArgs.size(), loopStepExpr.id())); } - CelCall newMacroCall = newMacroCallExpr.call(); - newMacroCallExpr.setCall( - newMacroCallExpr.call().toBuilder() - .clearArgs() - .addArgs( - newMacroCall.args().get(0)) // iter_var is first argument of the call by convention - .addArgs(args.get(1).createList().elements()) - .build()); + CelMutableCall existingMacroCall = newMacroCallExpr.call(); + CelMutableCall newMacroCall = + existingMacroCall.target().isPresent() + ? CelMutableCall.create(existingMacroCall.target().get(), existingMacroCall.function()) + : CelMutableCall.create(existingMacroCall.function()); + newMacroCall.addArgs( + existingMacroCall.args().get(0)); // iter_var is first argument of the call by convention + newMacroCall.addArgs(loopStepArgs.get(1).createList().elements()); + + newMacroCallExpr.setCall(newMacroCall); } - private CelExpr.Builder mutateExpr( + private CelMutableExpr mutateExpr( ExprIdGenerator idGenerator, - CelExpr.Builder root, - CelExpr.Builder newExpr, + CelMutableExpr root, + CelMutableExpr newExpr, long exprIdToReplace) { - MutableExprVisitor astMutator = + MutableExprVisitor mutableAst = MutableExprVisitor.newInstance(idGenerator, newExpr, exprIdToReplace, iterationLimit); - return astMutator.visit(root); + return mutableAst.visit(root); } - private CelExpr.Builder renumberExprIds(ExprIdGenerator idGenerator, CelExpr.Builder root) { - MutableExprVisitor astMutator = + private CelMutableExpr renumberExprIds(ExprIdGenerator idGenerator, CelMutableExpr root) { + MutableExprVisitor mutableAst = MutableExprVisitor.newInstance(idGenerator, root, Integer.MIN_VALUE, iterationLimit); - return astMutator.visit(root); + return mutableAst.visit(root); } - private static long getMaxId(CelAbstractSyntaxTree ast) { - long maxId = getMaxId(ast.getExpr()); - for (Entry macroCall : ast.getSource().getMacroCalls().entrySet()) { + private static long getMaxId(CelExpr newExpr) { + return CelNavigableExpr.fromExpr(newExpr) + .allNodes() + .mapToLong(CelNavigableExpr::id) + .max() + .orElseThrow(NoSuchElementException::new); + } + + private static long getMaxId(CelMutableAst mutableAst) { + return getMaxId(CelNavigableMutableAst.fromAst(mutableAst)); + } + + private static long getMaxId(CelNavigableMutableAst navAst) { + long maxId = navAst.getRoot().maxId(); + for (Entry macroCall : navAst.getAst().source().getMacroCalls().entrySet()) { maxId = max(maxId, getMaxId(macroCall.getValue())); } return maxId; } - private static long getMaxId(CelExpr newExpr) { - return CelNavigableExpr.fromExpr(newExpr) + private static long getMaxId(CelMutableExpr mutableExpr) { + return CelNavigableMutableExpr.fromExpr(mutableExpr) .allNodes() - .mapToLong(CelNavigableExpr::id) + .mapToLong(CelNavigableMutableExpr::id) .max() .orElseThrow(NoSuchElementException::new); } - private static int countComprehensionNestingLevel(CelNavigableExpr comprehensionExpr) { + private static int countComprehensionNestingLevel(CelNavigableMutableExpr comprehensionExpr) { int nestedLevel = 0; - Optional maybeParent = comprehensionExpr.parent(); + Optional maybeParent = comprehensionExpr.parent(); while (maybeParent.isPresent()) { if (maybeParent.get().getKind().equals(Kind.COMPREHENSION)) { nestedLevel++; @@ -822,14 +943,24 @@ private static int countComprehensionNestingLevel(CelNavigableExpr comprehension public abstract static class MangledComprehensionAst { /** AST after the iteration variables have been mangled. */ - public abstract CelAbstractSyntaxTree ast(); + public abstract CelMutableAst mutableAst(); /** Map containing the mangled identifier names to their types. */ public abstract ImmutableMap mangledComprehensionMap(); + /** + * TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. + * + * @deprecated Use MutableExpr based APIs instead. + */ + @Deprecated + public CelAbstractSyntaxTree ast() { + return mutableAst().toParsedAst(); + } + private static MangledComprehensionAst of( - CelAbstractSyntaxTree ast, + CelMutableAst ast, ImmutableMap mangledComprehensionMap) { return new AutoValue_AstMutator_MangledComprehensionAst(ast, mangledComprehensionMap); } @@ -846,11 +977,9 @@ public abstract static class MangledComprehensionType { public abstract Optional iterVarType(); /** Type of comprehension result */ - public abstract Optional resultType(); + public abstract CelType resultType(); - private static MangledComprehensionType of( - Optional iterVarType, Optional resultType) { - Preconditions.checkArgument(iterVarType.isPresent() || resultType.isPresent()); + private static MangledComprehensionType of(Optional iterVarType, CelType resultType) { return new AutoValue_AstMutator_MangledComprehensionType(iterVarType, resultType); } } @@ -872,24 +1001,4 @@ private static MangledComprehensionName of(String iterVarName, String resultName return new AutoValue_AstMutator_MangledComprehensionName(iterVarName, resultName); } } - - /** - * Intermediate value class to store the generated CelExpr for the bind macro and the macro call - * information. - */ - @AutoValue - abstract static class BindMacro { - /** Comprehension expr for the generated cel.bind macro. */ - abstract CelExpr bindExpr(); - - /** - * Call expr representation that will be stored in the macro call map of the AST. This is - * typically used for the purposes of supporting unparse. - */ - abstract CelExpr bindMacro(); - - private static BindMacro of(CelExpr bindExpr, CelExpr bindMacro) { - return new AutoValue_AstMutator_BindMacro(bindExpr, bindMacro); - } - } } diff --git a/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel index a6751d8a..d7833098 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel @@ -93,6 +93,7 @@ java_library( "//common/ast:mutable_ast", "//common/navigation", "//common/navigation:common", + "//common/navigation:mutable_navigation", "//common/types:type_providers", "@maven//:com_google_errorprone_error_prone_annotations", "@maven//:com_google_guava_guava", 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 76816810..7397d78d 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java +++ b/optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java @@ -222,12 +222,7 @@ private OptimizationResult optimizeUsingCelBlock(CelNavigableAst navigableAst, C iterVarType -> newVarDecls.add( CelVarDecl.newVarDeclaration(name.iterVarName(), iterVarType))); - type.resultType() - .ifPresent( - comprehensionResultType -> - newVarDecls.add( - CelVarDecl.newVarDeclaration( - name.resultName(), comprehensionResultType))); + newVarDecls.add(CelVarDecl.newVarDeclaration(name.resultName(), type.resultType())); }); // Type-check all sub-expressions then create new block index identifiers. diff --git a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java index 6eb0e04b..5f254dc0 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java +++ b/optimizer/src/test/java/dev/cel/optimizer/AstMutatorTest.java @@ -18,7 +18,6 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; -import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.testing.junit.testparameterinjector.TestParameterInjector; import com.google.testing.junit.testparameterinjector.TestParameters; @@ -33,12 +32,16 @@ import dev.cel.common.CelSource.Extension.Version; import dev.cel.common.ast.CelConstant; import dev.cel.common.ast.CelExpr; -import dev.cel.common.ast.CelExpr.CelCall; -import dev.cel.common.ast.CelExpr.CelIdent; -import dev.cel.common.ast.CelExpr.CelSelect; import dev.cel.common.ast.CelExpr.ExprKind.Kind; +import dev.cel.common.ast.CelMutableAst; +import dev.cel.common.ast.CelMutableExpr; +import dev.cel.common.ast.CelMutableExpr.CelMutableCall; +import dev.cel.common.ast.CelMutableExpr.CelMutableCreateList; +import dev.cel.common.ast.CelMutableExpr.CelMutableSelect; +import dev.cel.common.ast.CelMutableExprConverter; import dev.cel.common.navigation.CelNavigableAst; import dev.cel.common.navigation.CelNavigableExpr; +import dev.cel.common.navigation.CelNavigableMutableExpr; import dev.cel.common.types.SimpleType; import dev.cel.common.types.StructTypeReference; import dev.cel.extensions.CelExtensions; @@ -67,39 +70,44 @@ public class AstMutatorTest { .build(); private static final CelUnparser CEL_UNPARSER = CelUnparserFactory.newUnparser(); - private static final AstMutator MUTABLE_AST = AstMutator.newInstance(1000); + private static final AstMutator AST_MUTATOR = AstMutator.newInstance(1000); @Test - public void constExpr() throws Exception { + public void replaceSubtree_replaceConst() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("10").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newBooleanConst = CelMutableExpr.ofConstant(1, CelConstant.ofValue(true)); - CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(true)).build(), 1); + CelMutableAst result = + AST_MUTATOR.replaceSubtree(mutableAst, newBooleanConst, mutableAst.expr().id()); - assertThat(mutatedAst.getExpr()) + assertThat(result.toParsedAst().getExpr()) .isEqualTo(CelExpr.ofConstantExpr(3, CelConstant.ofValue(true))); } @Test public void astMutator_returnsParsedAst() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("10").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newBooleanConst = CelMutableExpr.ofConstant(1, CelConstant.ofValue(true)); - CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(true)).build(), 1); + CelMutableAst result = + AST_MUTATOR.replaceSubtree(mutableAst, newBooleanConst, mutableAst.expr().id()); assertThat(ast.isChecked()).isTrue(); - assertThat(mutatedAst.isChecked()).isFalse(); + assertThat(result.toParsedAst().isChecked()).isFalse(); } @Test public void astMutator_nonMacro_sourceCleared() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("10").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newBooleanConst = CelMutableExpr.ofConstant(1, CelConstant.ofValue(true)); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(true)).build(), 1); + AST_MUTATOR + .replaceSubtree(mutableAst, newBooleanConst, mutableAst.expr().id()) + .toParsedAst(); assertThat(mutatedAst.getSource().getDescription()).isEmpty(); assertThat(mutatedAst.getSource().getLineOffsets()).isEmpty(); @@ -111,10 +119,11 @@ public void astMutator_nonMacro_sourceCleared() throws Exception { @Test public void astMutator_macro_sourceMacroCallsPopulated() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("has(TestAllTypes{}.single_int32)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newBooleanConst = CelMutableExpr.ofConstant(1, CelConstant.ofValue(true)); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(true)).build(), 1); + AST_MUTATOR.replaceSubtree(mutableAst, newBooleanConst, 1).toParsedAst(); // no_op assertThat(mutatedAst.getSource().getDescription()).isEmpty(); assertThat(mutatedAst.getSource().getLineOffsets()).isEmpty(); @@ -131,10 +140,12 @@ public void replaceSubtree_astContainsTaggedExtension_retained() throws Exceptio ast = CelAbstractSyntaxTree.newCheckedAst( ast.getExpr(), celSource, ast.getReferenceMap(), ast.getTypeMap()); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(true)).build(), 1); + AST_MUTATOR + .replaceSubtree(mutableAst, CelMutableExpr.ofConstant(CelConstant.ofValue(true)), 1) + .toParsedAst(); assertThat(mutatedAst.getSource().getExtensions()).containsExactly(extension); } @@ -150,6 +161,7 @@ public void replaceSubtreeWithNewAst_astsContainTaggedExtension_retained() throw ast.getSource().toBuilder().addAllExtensions(extension).build(), ast.getReferenceMap(), ast.getTypeMap()); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); // Setup second AST with another test extension CelAbstractSyntaxTree astToReplaceWith = CEL.compile("cel.bind(a, true, a)").getAst(); Extension extension2 = Extension.create("test2", Version.of(2, 2)); @@ -159,10 +171,11 @@ public void replaceSubtreeWithNewAst_astsContainTaggedExtension_retained() throw astToReplaceWith.getSource().toBuilder().addAllExtensions(extension2).build(), astToReplaceWith.getReferenceMap(), astToReplaceWith.getTypeMap()); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(astToReplaceWith); // Mutate the original AST with the new AST at the root CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewAst(ast, astToReplaceWith, ast.getExpr().id()); + AST_MUTATOR.replaceSubtree(mutableAst, mutableAst2, mutableAst.expr().id()).toParsedAst(); // Expect that both the extensions are merged assertThat(mutatedAst.getSource().getExtensions()).containsExactly(extension, extension2); @@ -179,6 +192,7 @@ public void replaceSubtreeWithNewAst_astsContainSameExtensions_deduped() throws ast.getSource().toBuilder().addAllExtensions(extension).build(), ast.getReferenceMap(), ast.getTypeMap()); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); // Setup second AST with the same test extension as above CelAbstractSyntaxTree astToReplaceWith = CEL.compile("cel.bind(a, true, a)").getAst(); Extension extension2 = Extension.create("test", Version.of(1, 1)); @@ -188,10 +202,11 @@ public void replaceSubtreeWithNewAst_astsContainSameExtensions_deduped() throws astToReplaceWith.getSource().toBuilder().addAllExtensions(extension2).build(), astToReplaceWith.getReferenceMap(), astToReplaceWith.getTypeMap()); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(astToReplaceWith); // Mutate the original AST with the new AST at the root CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewAst(ast, astToReplaceWith, ast.getExpr().id()); + AST_MUTATOR.replaceSubtree(mutableAst, mutableAst2, mutableAst.expr().id()).toParsedAst(); // Expect that the extension is deduped assertThat(mutatedAst.getSource().getExtensions()).containsExactly(extension); @@ -207,10 +222,11 @@ public void replaceSubtree_rootReplacedWithMacro_macroCallPopulated( String source, int expectedMacroCallSize) throws Exception { CelAbstractSyntaxTree ast = CEL.compile("1").getAst(); CelAbstractSyntaxTree ast2 = CEL.compile(source).getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast2); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewAst( - ast, ast2, CelNavigableAst.fromAst(ast).getRoot().id()); + AST_MUTATOR.replaceSubtree(mutableAst, mutableAst2, mutableAst.expr().id()).toParsedAst(); assertThat(mutatedAst.getSource().getMacroCalls()).hasSize(expectedMacroCallSize); assertThat(CEL_UNPARSER.unparse(mutatedAst)).isEqualTo(source); @@ -218,30 +234,50 @@ public void replaceSubtree_rootReplacedWithMacro_macroCallPopulated( } @Test - public void replaceSubtree_branchReplacedWithMacro_macroCallPopulated() throws Exception { + public void replaceSubtree_leftBranchReplacedWithMacro_macroCallPopulated() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("true && false").getAst(); CelAbstractSyntaxTree ast2 = CEL.compile("[1].exists(x, x > 0)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast2); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewAst(ast, ast2, 3); // Replace false with the macro expr - CelAbstractSyntaxTree mutatedAst2 = - MUTABLE_AST.replaceSubtreeWithNewAst(ast, ast2, 1); // Replace true with the macro expr + AST_MUTATOR + .replaceSubtree(mutableAst, mutableAst2, 3) + .toParsedAst(); // Replace false with the macro expr assertThat(mutatedAst.getSource().getMacroCalls()).hasSize(1); assertThat(CEL_UNPARSER.unparse(mutatedAst)).isEqualTo("true && [1].exists(x, x > 0)"); assertThat(CEL.createProgram(CEL.check(mutatedAst).getAst()).eval()).isEqualTo(true); - assertThat(mutatedAst2.getSource().getMacroCalls()).hasSize(1); - assertThat(CEL_UNPARSER.unparse(mutatedAst2)).isEqualTo("[1].exists(x, x > 0) && false"); - assertThat(CEL.createProgram(CEL.check(mutatedAst2).getAst()).eval()).isEqualTo(false); + } + + @Test + public void replaceSubtree_rightBranchReplacedWithMacro_macroCallPopulated() throws Exception { + CelAbstractSyntaxTree ast = CEL.compile("true && false").getAst(); + CelAbstractSyntaxTree ast2 = CEL.compile("[1].exists(x, x > 0)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast2); + + CelAbstractSyntaxTree mutatedAst = + AST_MUTATOR + .replaceSubtree(mutableAst, mutableAst2, 1) + .toParsedAst(); // Replace true with the macro expr + + assertThat(mutatedAst.getSource().getMacroCalls()).hasSize(1); + assertThat(CEL_UNPARSER.unparse(mutatedAst)).isEqualTo("[1].exists(x, x > 0) && false"); + assertThat(CEL.createProgram(CEL.check(mutatedAst).getAst()).eval()).isEqualTo(false); } @Test public void replaceSubtree_macroInsertedIntoExistingMacro_macroCallPopulated() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("[1].exists(x, x > 0 && true)").getAst(); CelAbstractSyntaxTree ast2 = CEL.compile("[2].exists(y, y > 0)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast2); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewAst(ast, ast2, 9); // Replace true with the ast2 maro expr + AST_MUTATOR + .replaceSubtree(mutableAst, mutableAst2, 9) + .toParsedAst(); // Replace true with the ast2 maro expr assertThat(mutatedAst.getSource().getMacroCalls()).hasSize(2); assertThat(CEL_UNPARSER.unparse(mutatedAst)) @@ -252,24 +288,25 @@ public void replaceSubtree_macroInsertedIntoExistingMacro_macroCallPopulated() t @Test public void replaceSubtreeWithNewBindMacro_replaceRoot() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("1 + 1").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); String variableName = "@r0"; - CelExpr resultExpr = - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs( - CelExpr.ofIdentExpr(0, variableName), CelExpr.ofIdentExpr(0, variableName)) - .build()) - .build(); + CelMutableExpr resultExpr = + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + CelMutableExpr.ofIdent(variableName), + CelMutableExpr.ofIdent(variableName))); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewBindMacro( - ast, - variableName, - CelExpr.ofConstantExpr(0, CelConstant.ofValue(3L)), - resultExpr, - CelNavigableAst.fromAst(ast).getRoot().id()); + AST_MUTATOR + .replaceSubtreeWithNewBindMacro( + mutableAst, + variableName, + CelMutableExpr.ofConstant(CelConstant.ofValue(3L)), + resultExpr, + mutableAst.expr().id(), + true) + .toParsedAst(); assertThat(mutatedAst.getSource().getMacroCalls()).hasSize(1); assertThat(CEL_UNPARSER.unparse(mutatedAst)).isEqualTo("cel.bind(@r0, 3, @r0 + @r0)"); @@ -282,58 +319,45 @@ public void replaceSubtreeWithNewBindMacro_nestedBindMacro_replaceComprehensionR throws Exception { // Arrange CelAbstractSyntaxTree ast = CEL.compile("1 + 1").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); String variableName = "@r0"; - CelExpr resultExpr = - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs( - CelExpr.ofIdentExpr(0, variableName), CelExpr.ofIdentExpr(0, variableName)) - .build()) - .build(); + CelMutableExpr resultExpr = + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + CelMutableExpr.ofIdent(variableName), + CelMutableExpr.ofIdent(variableName))); // Act // Perform the initial replacement. (1 + 1) -> cel.bind(@r0, 3, @r0 + @r0) - CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewBindMacro( - ast, + mutableAst = + AST_MUTATOR.replaceSubtreeWithNewBindMacro( + mutableAst, variableName, - CelExpr.ofConstantExpr(0, CelConstant.ofValue(3L)), + CelMutableExpr.ofConstant(CelConstant.ofValue(3L)), resultExpr, - 2); // Replace + + 2, + true); // Replace + String nestedVariableName = "@r1"; // Construct a new result expression of the form @r0 + @r0 + @r1 + @r1 resultExpr = - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs( - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs( - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs( - CelExpr.ofIdentExpr(0, variableName), - CelExpr.ofIdentExpr(0, variableName)) - .build()) - .build(), - CelExpr.ofIdentExpr(0, nestedVariableName)) - .build()) - .build(), - CelExpr.ofIdentExpr(0, nestedVariableName)) - .build()) - .build(); + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + CelMutableExpr.ofIdent(variableName), + CelMutableExpr.ofIdent(variableName))), + CelMutableExpr.ofIdent(nestedVariableName))), + CelMutableExpr.ofIdent(nestedVariableName))); + // Find the call node (_+_) in the comprehension's result long exprIdToReplace = - CelNavigableAst.fromAst(mutatedAst) - .getRoot() + CelNavigableMutableExpr.fromExpr(mutableAst.expr()) .children() .filter( node -> @@ -344,14 +368,16 @@ public void replaceSubtreeWithNewBindMacro_nestedBindMacro_replaceComprehensionR .expr() .id(); // This should produce cel.bind(@r1, 1, cel.bind(@r0, 3, @r0 + @r0 + @r1 + @r1)) - mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewBindMacro( - mutatedAst, + mutableAst = + AST_MUTATOR.replaceSubtreeWithNewBindMacro( + mutableAst, nestedVariableName, - CelExpr.ofConstantExpr(0, CelConstant.ofValue(1L)), + CelMutableExpr.ofConstant(CelConstant.ofValue(1L)), resultExpr, - exprIdToReplace); // Replace + + exprIdToReplace, + true); // Replace + + CelAbstractSyntaxTree mutatedAst = mutableAst.toParsedAst(); assertThat(mutatedAst.getSource().getMacroCalls()).hasSize(2); assertThat(CEL.createProgram(CEL.check(mutatedAst).getAst()).eval()).isEqualTo(8); assertThat(CEL_UNPARSER.unparse(mutatedAst)) @@ -363,62 +389,56 @@ public void replaceSubtreeWithNewBindMacro_nestedBindMacro_replaceComprehensionR public void replaceSubtreeWithNewBindMacro_replaceRootWithNestedBindMacro() throws Exception { // Arrange CelAbstractSyntaxTree ast = CEL.compile("1 + 1 + 3 + 3").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); String variableName = "@r0"; - CelExpr resultExpr = - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs( - CelExpr.ofIdentExpr(0, variableName), CelExpr.ofIdentExpr(0, variableName)) - .build()) - .build(); + CelMutableExpr resultExpr = + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + CelMutableExpr.ofIdent(variableName), + CelMutableExpr.ofIdent(variableName))); // Act // Perform the initial replacement. (1 + 1 + 3 + 3) -> cel.bind(@r0, 1, @r0 + @r0) + 3 + 3 - CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewBindMacro( - ast, + mutableAst = + AST_MUTATOR.replaceSubtreeWithNewBindMacro( + mutableAst, variableName, - CelExpr.ofConstantExpr(0, CelConstant.ofValue(1L)), + CelMutableExpr.ofConstant(CelConstant.ofValue(1L)), resultExpr, - 2); // Replace + + 2, + true); // Replace + // Construct a new result expression of the form: // cel.bind(@r1, 3, cel.bind(@r0, 1, @r0 + @r0) + @r1 + @r1) String nestedVariableName = "@r1"; - CelExpr bindMacro = - CelNavigableAst.fromAst(mutatedAst) - .getRoot() + CelMutableExpr bindMacro = + CelNavigableMutableExpr.fromExpr(mutableAst.expr()) .descendants() .filter(node -> node.getKind().equals(Kind.COMPREHENSION)) .findAny() .get() .expr(); resultExpr = - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs( - CelExpr.newBuilder() - .setCall( - CelCall.newBuilder() - .setFunction(Operator.ADD.getFunction()) - .addArgs(bindMacro, CelExpr.ofIdentExpr(0, nestedVariableName)) - .build()) - .build(), - CelExpr.ofIdentExpr(0, nestedVariableName)) - .build()) - .build(); + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + CelMutableExpr.ofCall( + CelMutableCall.create( + Operator.ADD.getFunction(), + bindMacro, + CelMutableExpr.ofIdent(nestedVariableName))), + CelMutableExpr.ofIdent(nestedVariableName))); // Replace the root with the new result and a bind macro inserted - mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewBindMacro( - mutatedAst, + mutableAst = + AST_MUTATOR.replaceSubtreeWithNewBindMacro( + mutableAst, nestedVariableName, - CelExpr.ofConstantExpr(0, CelConstant.ofValue(3L)), + CelMutableExpr.ofConstant(CelConstant.ofValue(3L)), resultExpr, - CelNavigableAst.fromAst(mutatedAst).getRoot().id()); + mutableAst.expr().id(), + true); + CelAbstractSyntaxTree mutatedAst = mutableAst.toParsedAst(); assertThat(mutatedAst.getSource().getMacroCalls()).hasSize(2); assertThat(CEL.createProgram(CEL.check(mutatedAst).getAst()).eval()).isEqualTo(8); assertThat(CEL_UNPARSER.unparse(mutatedAst)) @@ -431,10 +451,11 @@ public void replaceSubtree_macroReplacedWithConstExpr_macroCallCleared() throws CelAbstractSyntaxTree ast = CEL.compile("[1].exists(x, x > 0) && [2].exists(x, x > 0)").getAst(); CelAbstractSyntaxTree ast2 = CEL.compile("1").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast2); CelAbstractSyntaxTree mutatedAst = - MUTABLE_AST.replaceSubtreeWithNewAst( - ast, ast2, CelNavigableAst.fromAst(ast).getRoot().id()); + AST_MUTATOR.replaceSubtree(mutableAst, mutableAst2, mutableAst.expr().id()).toParsedAst(); assertThat(mutatedAst.getSource().getMacroCalls()).isEmpty(); assertThat(CEL_UNPARSER.unparse(mutatedAst)).isEqualTo("1"); @@ -461,19 +482,23 @@ public void replaceSubtree_replaceExtraneousListCreatedByMacro_unparseSuccess() // } // } CelAbstractSyntaxTree ast = CEL.compile("[1].map(x, 1)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableAst mutableAst2 = CelMutableAst.fromCelAst(ast); // These two mutation are equivalent. CelAbstractSyntaxTree mutatedAstWithList = - MUTABLE_AST.replaceSubtree( - ast, - CelExpr.ofCreateListExpr( - 0, - ImmutableList.of(CelExpr.newBuilder().setConstant(CelConstant.ofValue(2L)).build()), - ImmutableList.of()), - 9L); + AST_MUTATOR + .replaceSubtree( + mutableAst, + CelMutableExpr.ofCreateList( + CelMutableCreateList.create( + CelMutableExpr.ofConstant(CelConstant.ofValue(2L)))), + 9L) + .toParsedAst(); CelAbstractSyntaxTree mutatedAstWithConstant = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(2L)).build(), 5L); + AST_MUTATOR + .replaceSubtree(mutableAst2, CelMutableExpr.ofConstant(CelConstant.ofValue(2L)), 5L) + .toParsedAst(); assertThat(CEL_UNPARSER.unparse(mutatedAstWithList)).isEqualTo("[1].map(x, 2)"); assertThat(CEL_UNPARSER.unparse(mutatedAstWithConstant)).isEqualTo("[1].map(x, 2)"); @@ -488,12 +513,13 @@ public void globalCallExpr_replaceRoot() throws Exception { // + [2] x [5] // 1 [1] 2 [3] CelAbstractSyntaxTree ast = CEL.compile("1 + 2 + x").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(10)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(10)).build(), 4); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, mutableAst.expr().id()); - assertThat(replacedAst.getExpr()).isEqualTo(CelExpr.ofConstantExpr(7, CelConstant.ofValue(10))); + assertThat(result.toParsedAst().getExpr()) + .isEqualTo(CelExpr.ofConstantExpr(7, CelConstant.ofValue(10))); } @Test @@ -503,12 +529,12 @@ public void globalCallExpr_replaceLeaf() throws Exception { // + [2] x [5] // 1 [1] 2 [3] CelAbstractSyntaxTree ast = CEL.compile("1 + 2 + x").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(10)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(10)).build(), 1); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 1); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("10 + 2 + x"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10 + 2 + x"); } @Test @@ -518,12 +544,12 @@ public void globalCallExpr_replaceMiddleBranch() throws Exception { // + [2] x [5] // 1 [1] 2 [3] CelAbstractSyntaxTree ast = CEL.compile("1 + 2 + x").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(10)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(10)).build(), 2); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 2); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("10 + x"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10 + x"); } @Test @@ -534,10 +560,12 @@ public void globalCallExpr_replaceMiddleBranch_withCallExpr() throws Exception { // 1 [1] 2 [3] CelAbstractSyntaxTree ast = CEL.compile("1 + 2 + x").getAst(); CelAbstractSyntaxTree ast2 = CEL.compile("4 + 5 + 6").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExprConverter.fromCelExpr(ast2.getExpr()); - CelAbstractSyntaxTree replacedAst = MUTABLE_AST.replaceSubtree(ast, ast2.getExpr(), 2); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 2); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("4 + 5 + 6 + x"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("4 + 5 + 6 + x"); } @Test @@ -555,12 +583,12 @@ public void memberCallExpr_replaceLeafTarget() throws Exception { "func_overload", SimpleType.INT, SimpleType.INT, SimpleType.INT))) .build(); CelAbstractSyntaxTree ast = cel.compile("10.func(4.func(5))").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(20)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(20)).build(), 3); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 3); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("10.func(20.func(5))"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10.func(20.func(5))"); } @Test @@ -578,12 +606,12 @@ public void memberCallExpr_replaceLeafArgument() throws Exception { "func_overload", SimpleType.INT, SimpleType.INT, SimpleType.INT))) .build(); CelAbstractSyntaxTree ast = cel.compile("10.func(4.func(5))").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(20)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(20)).build(), 5); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 5); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("10.func(4.func(20))"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10.func(4.func(20))"); } @Test @@ -601,12 +629,12 @@ public void memberCallExpr_replaceMiddleBranchTarget() throws Exception { "func_overload", SimpleType.INT, SimpleType.INT, SimpleType.INT))) .build(); CelAbstractSyntaxTree ast = cel.compile("10.func(4.func(5))").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(20)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(20)).build(), 1); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 1); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("20.func(4.func(5))"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("20.func(4.func(5))"); } @Test @@ -624,12 +652,12 @@ public void memberCallExpr_replaceMiddleBranchArgument() throws Exception { "func_overload", SimpleType.INT, SimpleType.INT, SimpleType.INT))) .build(); CelAbstractSyntaxTree ast = cel.compile("10.func(4.func(5))").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(20)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(20)).build(), 4); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 4); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("10.func(20)"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("10.func(20)"); } @Test @@ -639,23 +667,14 @@ public void select_replaceField() throws Exception { // 5 [1] select [4] // msg [3] CelAbstractSyntaxTree ast = CEL.compile("5 + msg.single_int64").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = + CelMutableExpr.ofSelect( + CelMutableSelect.create(CelMutableExpr.ofIdent("test"), "single_sint32")); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, - CelExpr.newBuilder() - .setSelect( - CelSelect.newBuilder() - .setField("single_sint32") - .setOperand( - CelExpr.newBuilder() - .setIdent(CelIdent.newBuilder().setName("test").build()) - .build()) - .build()) - .build(), - 4); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 4); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("5 + test.single_sint32"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("5 + test.single_sint32"); } @Test @@ -665,14 +684,12 @@ public void select_replaceOperand() throws Exception { // 5 [1] select [4] // msg [3] CelAbstractSyntaxTree ast = CEL.compile("5 + msg.single_int64").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofIdent("test"); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, - CelExpr.newBuilder().setIdent(CelIdent.newBuilder().setName("test").build()).build(), - 3); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 3); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("5 + test.single_int64"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("5 + test.single_int64"); } @Test @@ -681,12 +698,12 @@ public void list_replaceElement() throws Exception { // list [1] // 2 [2] 3 [3] 4 [4] CelAbstractSyntaxTree ast = CEL.compile("[2, 3, 4]").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(5)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(5)).build(), 4); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 4); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("[2, 3, 5]"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("[2, 3, 5]"); } @Test @@ -696,12 +713,13 @@ public void createStruct_replaceValue() throws Exception { // single_int64 [2] // 2 [3] CelAbstractSyntaxTree ast = CEL.compile("TestAllTypes{single_int64: 2}").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(5)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(5)).build(), 3); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 3); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("TestAllTypes{single_int64: 5}"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())) + .isEqualTo("TestAllTypes{single_int64: 5}"); } @Test @@ -711,12 +729,12 @@ public void createMap_replaceKey() throws Exception { // map_entry [2] // 'a' [3] : 1 [4] CelAbstractSyntaxTree ast = CEL.compile("{'a': 1}").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(5)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(5)).build(), 3); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 3); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("{5: 1}"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("{5: 1}"); } @Test @@ -726,34 +744,36 @@ public void createMap_replaceValue() throws Exception { // map_entry [2] // 'a' [3] : 1 [4] CelAbstractSyntaxTree ast = CEL.compile("{'a': 1}").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(5)); - CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(5)).build(), 4); + CelMutableAst result = AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 4); - assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("{\"a\": 5}"); + assertThat(CEL_UNPARSER.unparse(result.toParsedAst())).isEqualTo("{\"a\": 5}"); } @Test public void comprehension_replaceIterRange() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("[true].exists(i, i)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(false)); CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(false)).build(), 2); + AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 2).toParsedAst(); assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("[false].exists(i, i)"); assertThat(CEL.createProgram(CEL.check(replacedAst).getAst()).eval()).isEqualTo(false); - assertConsistentMacroCalls(ast); + assertConsistentMacroCalls(replacedAst); } @Test public void comprehension_replaceAccuInit() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("[false].exists(i, i)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(true)); CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, CelExpr.newBuilder().setConstant(CelConstant.ofValue(true)).build(), 6); + AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 6).toParsedAst(); assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("[false].exists(i, i)"); assertThat(CEL.createProgram(CEL.check(replacedAst).getAst()).eval()).isEqualTo(true); @@ -766,12 +786,11 @@ public void comprehension_replaceAccuInit() throws Exception { @Test public void comprehension_replaceLoopStep() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("[false].exists(i, i)").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); + CelMutableExpr newExpr = CelMutableExpr.ofIdent("test"); CelAbstractSyntaxTree replacedAst = - MUTABLE_AST.replaceSubtree( - ast, - CelExpr.newBuilder().setIdent(CelIdent.newBuilder().setName("test").build()).build(), - 5); + AST_MUTATOR.replaceSubtree(mutableAst, newExpr, 5).toParsedAst(); assertThat(CEL_UNPARSER.unparse(replacedAst)).isEqualTo("[false].exists(i, test)"); assertConsistentMacroCalls(ast); @@ -782,7 +801,7 @@ public void mangleComprehensionVariable_singleMacro() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("[false].exists(i, i)").getAst(); CelAbstractSyntaxTree mangledAst = - MUTABLE_AST.mangleComprehensionIdentifierNames(ast, "@c", "@x").ast(); + AST_MUTATOR.mangleComprehensionIdentifierNames(ast, "@c", "@x").mutableAst().toParsedAst(); assertThat(mangledAst.getExpr().toString()) .isEqualTo( @@ -838,12 +857,28 @@ public void mangleComprehensionVariable_singleMacro() throws Exception { assertConsistentMacroCalls(ast); } + @Test + public void mangleComprehensionVariable_macroSourceDisabled_macroCallMapIsEmpty() + throws Exception { + Cel cel = + CelFactory.standardCelBuilder() + .setStandardMacros(CelStandardMacro.STANDARD_MACROS) + .setOptions(CelOptions.current().populateMacroCalls(false).build()) + .build(); + CelAbstractSyntaxTree ast = cel.compile("[false].exists(i, i)").getAst(); + + CelAbstractSyntaxTree mangledAst = + AST_MUTATOR.mangleComprehensionIdentifierNames(ast, "@c", "@x").mutableAst().toParsedAst(); + + assertThat(mangledAst.getSource().getMacroCalls()).isEmpty(); + } + @Test public void mangleComprehensionVariable_nestedMacroWithShadowedVariables() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("[x].exists(x, [x].exists(x, x == 1))").getAst(); CelAbstractSyntaxTree mangledAst = - MUTABLE_AST.mangleComprehensionIdentifierNames(ast, "@c", "@x").ast(); + AST_MUTATOR.mangleComprehensionIdentifierNames(ast, "@c", "@x").mutableAst().toParsedAst(); assertThat(mangledAst.getExpr().toString()) .isEqualTo( @@ -960,7 +995,7 @@ public void mangleComprehensionVariable_hasMacro_noOp() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("has(msg.single_int64)").getAst(); CelAbstractSyntaxTree mangledAst = - MUTABLE_AST.mangleComprehensionIdentifierNames(ast, "@c", "@x").ast(); + AST_MUTATOR.mangleComprehensionIdentifierNames(ast, "@c", "@x").mutableAst().toParsedAst(); assertThat(CEL_UNPARSER.unparse(mangledAst)).isEqualTo("has(msg.single_int64)"); assertThat( @@ -973,6 +1008,7 @@ public void mangleComprehensionVariable_hasMacro_noOp() throws Exception { @Test public void replaceSubtree_iterationLimitReached_throws() throws Exception { CelAbstractSyntaxTree ast = CEL.compile("true && false").getAst(); + CelMutableAst mutableAst = CelMutableAst.fromCelAst(ast); AstMutator astMutator = AstMutator.newInstance(1); IllegalStateException e = @@ -980,7 +1016,7 @@ public void replaceSubtree_iterationLimitReached_throws() throws Exception { IllegalStateException.class, () -> astMutator.replaceSubtree( - ast, CelExpr.ofConstantExpr(0, CelConstant.ofValue(false)), 1)); + mutableAst, CelMutableExpr.ofConstant(CelConstant.ofValue(false)), 1)); assertThat(e).hasMessageThat().isEqualTo("Max iteration count reached."); } diff --git a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel index 88e24415..4d751e4f 100644 --- a/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel +++ b/optimizer/src/test/java/dev/cel/optimizer/BUILD.bazel @@ -13,7 +13,9 @@ java_library( "//common:compiler_common", "//common:options", "//common/ast", + "//common/ast:mutable_ast", "//common/navigation", + "//common/navigation:mutable_navigation", "//common/resources/testdata/proto3:test_all_types_java_proto", "//common/types", "//compiler",