From 473e43942df8d3d3a57b47fd599325cde0fc31f3 Mon Sep 17 00:00:00 2001 From: Sokwhan Huh Date: Tue, 16 Apr 2024 16:14:39 -0700 Subject: [PATCH] Change MutableExprVisitor to visit using CelMutableExpr PiperOrigin-RevId: 625485309 --- .../main/java/dev/cel/optimizer/BUILD.bazel | 1 + .../dev/cel/optimizer/MutableExprVisitor.java | 142 ++++++++++-------- 2 files changed, 82 insertions(+), 61 deletions(-) diff --git a/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel b/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel index 7e2022ab1..07de0dd06 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel +++ b/optimizer/src/main/java/dev/cel/optimizer/BUILD.bazel @@ -90,6 +90,7 @@ java_library( "//common/annotations", "//common/ast", "//common/ast:expr_factory", + "//common/ast:mutable_ast", "//common/navigation", "//common/types:type_providers", "@maven//:com_google_errorprone_error_prone_annotations", diff --git a/optimizer/src/main/java/dev/cel/optimizer/MutableExprVisitor.java b/optimizer/src/main/java/dev/cel/optimizer/MutableExprVisitor.java index 36cb86069..ea751313a 100644 --- a/optimizer/src/main/java/dev/cel/optimizer/MutableExprVisitor.java +++ b/optimizer/src/main/java/dev/cel/optimizer/MutableExprVisitor.java @@ -15,146 +15,166 @@ package dev.cel.optimizer; import com.google.common.base.Preconditions; -import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; import dev.cel.common.annotations.Internal; 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.CelCreateList; -import dev.cel.common.ast.CelExpr.CelCreateMap; -import dev.cel.common.ast.CelExpr.CelCreateStruct; -import dev.cel.common.ast.CelExpr.CelSelect; import dev.cel.common.ast.CelExprIdGeneratorFactory.ExprIdGenerator; +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.CelMutableExpr.CelMutableCreateMap; +import dev.cel.common.ast.CelMutableExpr.CelMutableCreateStruct; +import dev.cel.common.ast.CelMutableExpr.CelMutableSelect; +import dev.cel.common.ast.CelMutableExprConverter; +import java.util.List; /** - * MutableExprVisitor performs mutation of {@link CelExpr} based on its configured parameters. + * MutableExprVisitor performs mutation of {@link CelMutableExpr} based on its configured + * parameters. * *

This class is NOT thread-safe. Callers should spawn a new instance of this class each time the * expression is being mutated. - * - *

Note that CelExpr is immutable by design. Therefore, the logic here doesn't actually mutate - * the existing expression tree. Instead, a brand new CelExpr is produced with the subtree swapped - * at the desired expression ID to replace. */ @Internal final class MutableExprVisitor { - private final CelExpr.Builder newExpr; + private final CelMutableExpr newExpr; private final ExprIdGenerator celExprIdGenerator; private final long iterationLimit; private int iterationCount; private long exprIdToReplace; + /** TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. */ static MutableExprVisitor newInstance( ExprIdGenerator idGenerator, CelExpr.Builder newExpr, long exprIdToReplace, long iterationLimit) { + return newInstance( + idGenerator, + CelMutableExprConverter.fromCelExpr(newExpr.build()), + exprIdToReplace, + iterationLimit); + } + + static MutableExprVisitor newInstance( + ExprIdGenerator idGenerator, + CelMutableExpr newExpr, + long exprIdToReplace, + long iterationLimit) { // iterationLimit * 2, because the expr can be walked twice due to the immutable nature of // CelExpr. return new MutableExprVisitor(idGenerator, newExpr, exprIdToReplace, iterationLimit * 2); } + /** TODO: Temporarily kept to retain method signature. Remove in the upcoming CLs. */ CelExpr.Builder visit(CelExpr.Builder root) { + CelMutableExpr mutatedRoot = visit(CelMutableExprConverter.fromCelExpr(root.build())); + return CelMutableExprConverter.fromMutableExpr(mutatedRoot).toBuilder(); + } + + CelMutableExpr visit(CelMutableExpr root) { if (++iterationCount > iterationLimit) { throw new IllegalStateException("Max iteration count reached."); } if (root.id() == exprIdToReplace) { exprIdToReplace = Integer.MIN_VALUE; // Marks that the subtree has been replaced. - return visit(this.newExpr.setId(root.id())); + this.newExpr.setId(root.id()); + return visit(this.newExpr); } root.setId(celExprIdGenerator.generate(root.id())); - switch (root.exprKind().getKind()) { + switch (root.getKind()) { case SELECT: - return visit(root, root.select().toBuilder()); + return visit(root, root.select()); case CALL: - return visit(root, root.call().toBuilder()); + return visit(root, root.call()); case CREATE_LIST: - return visit(root, root.createList().toBuilder()); + return visit(root, root.createList()); case CREATE_STRUCT: - return visit(root, root.createStruct().toBuilder()); + return visit(root, root.createStruct()); case CREATE_MAP: - return visit(root, root.createMap().toBuilder()); + return visit(root, root.createMap()); case COMPREHENSION: - return visit(root, root.comprehension().toBuilder()); + return visit(root, root.comprehension()); case CONSTANT: // Fall-through is intended case IDENT: case NOT_SET: // Note: comprehension arguments can contain a not set root. return root; } - throw new IllegalArgumentException("unexpected root kind: " + root.exprKind().getKind()); + throw new IllegalArgumentException("unexpected root kind: " + root.getKind()); } - private CelExpr.Builder visit(CelExpr.Builder expr, CelSelect.Builder select) { - select.setOperand(visit(select.operand().toBuilder()).build()); - return expr.setSelect(select.build()); + @CanIgnoreReturnValue + private CelMutableExpr visit(CelMutableExpr expr, CelMutableSelect select) { + select.setOperand(visit(select.operand())); + return expr; } - private CelExpr.Builder visit(CelExpr.Builder expr, CelCall.Builder call) { + @CanIgnoreReturnValue + private CelMutableExpr visit(CelMutableExpr expr, CelMutableCall call) { if (call.target().isPresent()) { - call.setTarget(visit(call.target().get().toBuilder()).build()); + call.setTarget(visit(call.target().get())); } - ImmutableList argsBuilders = call.getArgsBuilders(); + List argsBuilders = call.args(); for (int i = 0; i < argsBuilders.size(); i++) { - CelExpr.Builder arg = argsBuilders.get(i); - call.setArg(i, visit(arg).build()); + CelMutableExpr arg = argsBuilders.get(i); + call.setArg(i, visit(arg)); } - return expr.setCall(call.build()); + return expr; } - private CelExpr.Builder visit(CelExpr.Builder expr, CelCreateStruct.Builder createStruct) { - ImmutableList entries = createStruct.getEntriesBuilders(); - for (int i = 0; i < entries.size(); i++) { - CelCreateStruct.Entry.Builder entry = entries.get(i); + @CanIgnoreReturnValue + private CelMutableExpr visit(CelMutableExpr expr, CelMutableCreateStruct createStruct) { + List entries = createStruct.entries(); + for (CelMutableCreateStruct.Entry entry : entries) { entry.setId(celExprIdGenerator.generate(entry.id())); - entry.setValue(visit(entry.value().toBuilder()).build()); - - createStruct.setEntry(i, entry.build()); + entry.setValue(visit(entry.value())); } - return expr.setCreateStruct(createStruct.build()); + return expr; } - private CelExpr.Builder visit(CelExpr.Builder expr, CelCreateMap.Builder createMap) { - ImmutableList entriesBuilders = createMap.getEntriesBuilders(); - for (int i = 0; i < entriesBuilders.size(); i++) { - CelCreateMap.Entry.Builder entry = entriesBuilders.get(i); + @CanIgnoreReturnValue + private CelMutableExpr visit(CelMutableExpr expr, CelMutableCreateMap createMap) { + List entriesBuilders = createMap.entries(); + for (CelMutableCreateMap.Entry entry : entriesBuilders) { entry.setId(celExprIdGenerator.generate(entry.id())); - entry.setKey(visit(entry.key().toBuilder()).build()); - entry.setValue(visit(entry.value().toBuilder()).build()); - - createMap.setEntry(i, entry.build()); + entry.setKey(visit(entry.key())); + entry.setValue(visit(entry.value())); } - return expr.setCreateMap(createMap.build()); + return expr; } - private CelExpr.Builder visit(CelExpr.Builder expr, CelCreateList.Builder createList) { - ImmutableList elementsBuilders = createList.getElementsBuilders(); + @CanIgnoreReturnValue + private CelMutableExpr visit(CelMutableExpr expr, CelMutableCreateList createList) { + List elementsBuilders = createList.elements(); for (int i = 0; i < elementsBuilders.size(); i++) { - CelExpr.Builder elem = elementsBuilders.get(i); - createList.setElement(i, visit(elem).build()); + CelMutableExpr elem = elementsBuilders.get(i); + createList.setElement(i, visit(elem)); } - return expr.setCreateList(createList.build()); + return expr; } - private CelExpr.Builder visit(CelExpr.Builder expr, CelComprehension.Builder comprehension) { - comprehension.setIterRange(visit(comprehension.iterRange().toBuilder()).build()); - comprehension.setAccuInit(visit(comprehension.accuInit().toBuilder()).build()); - comprehension.setLoopCondition(visit(comprehension.loopCondition().toBuilder()).build()); - comprehension.setLoopStep(visit(comprehension.loopStep().toBuilder()).build()); - comprehension.setResult(visit(comprehension.result().toBuilder()).build()); + @CanIgnoreReturnValue + private CelMutableExpr visit(CelMutableExpr expr, CelMutableComprehension comprehension) { + comprehension.setIterRange(visit(comprehension.iterRange())); + comprehension.setAccuInit(visit(comprehension.accuInit())); + comprehension.setLoopCondition(visit(comprehension.loopCondition())); + comprehension.setLoopStep(visit(comprehension.loopStep())); + comprehension.setResult(visit(comprehension.result())); - return expr.setComprehension(comprehension.build()); + return expr; } private MutableExprVisitor( ExprIdGenerator celExprIdGenerator, - CelExpr.Builder newExpr, + CelMutableExpr newExpr, long exprId, long iterationLimit) { Preconditions.checkState(iterationLimit > 0L);