diff --git a/common/src/main/java/dev/cel/common/ast/BUILD.bazel b/common/src/main/java/dev/cel/common/ast/BUILD.bazel
index ac4df5487..9ed6c80fc 100644
--- a/common/src/main/java/dev/cel/common/ast/BUILD.bazel
+++ b/common/src/main/java/dev/cel/common/ast/BUILD.bazel
@@ -34,6 +34,7 @@ EXPR_FACTORY_SOURCES = [
# keep sorted
MUTABLE_AST_SOURCES = [
+ "CelMutableAst.java",
"CelMutableExpr.java",
"CelMutableExprConverter.java",
]
@@ -124,6 +125,8 @@ java_library(
],
deps = [
":ast",
+ "//common",
+ "//common/types:type_providers",
"@maven//:com_google_guava_guava",
],
)
diff --git a/common/src/main/java/dev/cel/common/ast/CelMutableAst.java b/common/src/main/java/dev/cel/common/ast/CelMutableAst.java
new file mode 100644
index 000000000..f529e2e8d
--- /dev/null
+++ b/common/src/main/java/dev/cel/common/ast/CelMutableAst.java
@@ -0,0 +1,109 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.common.ast;
+
+import dev.cel.common.CelAbstractSyntaxTree;
+import dev.cel.common.CelSource;
+import dev.cel.common.types.CelType;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+/**
+ * An abstract representation of CEL Abstract Syntax tree that allows mutation in any of its
+ * properties. This class is semantically the same as that of the immutable {@link
+ * CelAbstractSyntaxTree}.
+ *
+ *
This should only be used within optimizers to augment an AST.
+ */
+public final class CelMutableAst {
+ private final CelMutableExpr mutatedExpr;
+ private final CelSource.Builder source;
+ private final Map references;
+ private final Map types;
+
+ /** Returns the underlying {@link CelMutableExpr} representation of the abstract syntax tree. */
+ public CelMutableExpr expr() {
+ return mutatedExpr;
+ }
+
+ /**
+ * Returns the {@link CelSource} that was used during construction of the abstract syntax tree.
+ */
+ public CelSource.Builder source() {
+ return source;
+ }
+
+ /**
+ * Returns the resolved reference to a declaration at expression ID for a type-checked AST.
+ *
+ * @return Optional of {@link CelReference} or {@link Optional#empty} if the reference does not
+ * exist at the ID.
+ */
+ public Optional getReference(long exprId) {
+ return Optional.ofNullable(references.get(exprId));
+ }
+
+ /**
+ * Returns the type of the expression node for a type-checked AST.
+ *
+ * @return Optional of {@link CelType} or {@link Optional#empty} if the type does not exist at the
+ * ID.
+ */
+ public Optional getType(long exprId) {
+ return Optional.ofNullable(types.get(exprId));
+ }
+
+ /** Converts this mutable AST into a parsed {@link CelAbstractSyntaxTree}. */
+ public CelAbstractSyntaxTree toParsedAst() {
+ return CelAbstractSyntaxTree.newParsedAst(
+ CelMutableExprConverter.fromMutableExpr(mutatedExpr), source.build());
+ }
+
+ /**
+ * Constructs an instance of {@link CelMutableAst} with the incoming {@link
+ * CelAbstractSyntaxTree}.
+ */
+ public static CelMutableAst fromCelAst(CelAbstractSyntaxTree ast) {
+ return new CelMutableAst(
+ CelMutableExprConverter.fromCelExpr(ast.getExpr()),
+ ast.getSource().toBuilder(),
+ ast.getReferenceMap(),
+ ast.getTypeMap());
+ }
+
+ /**
+ * Constructs an instance of {@link CelMutableAst} with the mutable expression and its source
+ * builder.
+ */
+ public static CelMutableAst of(CelMutableExpr mutableExpr, CelSource.Builder sourceBuilder) {
+ return new CelMutableAst(mutableExpr, sourceBuilder);
+ }
+
+ private CelMutableAst(CelMutableExpr mutatedExpr, CelSource.Builder source) {
+ this(mutatedExpr, source, new HashMap<>(), new HashMap<>());
+ }
+
+ private CelMutableAst(
+ CelMutableExpr mutatedExpr,
+ CelSource.Builder source,
+ Map references,
+ Map types) {
+ this.mutatedExpr = mutatedExpr;
+ this.source = source;
+ this.references = new HashMap<>(references);
+ this.types = new HashMap<>(types);
+ }
+}
diff --git a/common/src/test/java/dev/cel/common/ast/BUILD.bazel b/common/src/test/java/dev/cel/common/ast/BUILD.bazel
index 4c74f4964..103f4804a 100644
--- a/common/src/test/java/dev/cel/common/ast/BUILD.bazel
+++ b/common/src/test/java/dev/cel/common/ast/BUILD.bazel
@@ -13,6 +13,7 @@ java_library(
"//:java_truth",
"//common",
"//common:compiler_common",
+ "//common:options",
"//common/ast",
"//common/ast:cel_expr_visitor",
"//common/ast:expr_converter",
diff --git a/common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java b/common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java
new file mode 100644
index 000000000..1b249484d
--- /dev/null
+++ b/common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java
@@ -0,0 +1,91 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dev.cel.common.ast;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import dev.cel.common.CelAbstractSyntaxTree;
+import dev.cel.common.CelOptions;
+import dev.cel.common.CelSource;
+import dev.cel.common.types.SimpleType;
+import dev.cel.compiler.CelCompiler;
+import dev.cel.compiler.CelCompilerFactory;
+import dev.cel.parser.CelStandardMacro;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public final class CelMutableAstTest {
+
+ @Test
+ public void constructMutableAst() {
+ CelMutableExpr mutableExpr = CelMutableExpr.ofConstant(1L, CelConstant.ofValue("hello world"));
+ CelSource.Builder sourceBuilder = CelSource.newBuilder();
+
+ CelMutableAst celMutableAst = CelMutableAst.of(mutableExpr, sourceBuilder);
+
+ assertThat(celMutableAst.expr()).isEqualTo(mutableExpr);
+ assertThat(celMutableAst.source()).isSameInstanceAs(sourceBuilder);
+ }
+
+ @Test
+ public void fromCelAst_mutableAst_containsMutableExpr() throws Exception {
+ CelCompiler celCompiler = CelCompilerFactory.standardCelCompilerBuilder().build();
+ CelAbstractSyntaxTree ast = celCompiler.compile("'hello world'").getAst();
+
+ CelMutableAst celMutableAst = CelMutableAst.fromCelAst(ast);
+
+ assertThat(celMutableAst.expr())
+ .isEqualTo(CelMutableExpr.ofConstant(1L, CelConstant.ofValue("hello world")));
+ }
+
+ @Test
+ public void getType_success() throws Exception {
+ CelCompiler celCompiler = CelCompilerFactory.standardCelCompilerBuilder().build();
+ CelAbstractSyntaxTree ast = celCompiler.compile("'hello world'").getAst();
+ CelMutableAst celMutableAst = CelMutableAst.fromCelAst(ast);
+
+ assertThat(celMutableAst.getType(1L)).hasValue(SimpleType.STRING);
+ }
+
+ @Test
+ public void getReference_success() throws Exception {
+ CelCompiler celCompiler = CelCompilerFactory.standardCelCompilerBuilder().build();
+ CelAbstractSyntaxTree ast = celCompiler.compile("size('test')").getAst();
+ CelMutableAst celMutableAst = CelMutableAst.fromCelAst(ast);
+
+ assertThat(celMutableAst.getReference(1L))
+ .hasValue(CelReference.newBuilder().addOverloadIds("size_string").build());
+ }
+
+ @Test
+ public void parsedAst_roundTrip() throws Exception {
+ CelCompiler celCompiler =
+ CelCompilerFactory.standardCelCompilerBuilder()
+ .setOptions(CelOptions.current().populateMacroCalls(true).build())
+ .setStandardMacros(CelStandardMacro.STANDARD_MACROS)
+ .build();
+ CelAbstractSyntaxTree ast = celCompiler.parse("[1].exists(x, x > 0)").getAst();
+ CelMutableAst celMutableAst = CelMutableAst.fromCelAst(ast);
+
+ CelAbstractSyntaxTree roundTrippedAst = celMutableAst.toParsedAst();
+
+ assertThat(ast.getExpr()).isEqualTo(roundTrippedAst.getExpr());
+ assertThat(roundTrippedAst.getSource().getMacroCalls()).hasSize(1);
+ assertThat(roundTrippedAst.getSource().getMacroCalls())
+ .containsExactlyEntriesIn(ast.getSource().getMacroCalls());
+ }
+}