Add CelMutableAst

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",
     ],
 )

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}.
+ *
+ * <p>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<Long, CelReference> references;
+  private final Map<Long, CelType> 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<CelReference> 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<CelType> 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<Long, CelReference> references,
+      Map<Long, CelType> types) {
+    this.mutatedExpr = mutatedExpr;
+    this.source = source;
+    this.references = new HashMap<>(references);
+    this.types = new HashMap<>(types);
+  }
+}

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",

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());
+  }
+}