Change AstMutator to augment macro source using mutable expr

common/BUILD.bazel

@@ -39,6 +39,16 @@ java_library(
     exports = ["//common/src/main/java/dev/cel/common:error_codes"],
 )
 
+java_library(
+    name = "mutable_ast",
+    exports = ["//common/src/main/java/dev/cel/common:mutable_ast"],
+)
+
+java_library(
+    name = "mutable_source",
+    exports = ["//common/src/main/java/dev/cel/common:mutable_source"],
+)
+
 java_library(
     name = "runtime_exception",
     visibility = ["//visibility:public"],

common/ast/BUILD.bazel

@@ -34,6 +34,6 @@ java_library(
 )
 
 java_library(
-    name = "mutable_ast",
-    exports = ["//common/src/main/java/dev/cel/common/ast:mutable_ast"],
+    name = "mutable_expr",
+    exports = ["//common/src/main/java/dev/cel/common/ast:mutable_expr"],
 )

common/src/main/java/dev/cel/common/BUILD.bazel

@@ -134,3 +134,31 @@ java_library(
         "//common/annotations",
     ],
 )
+
+java_library(
+    name = "mutable_ast",
+    srcs = ["CelMutableAst.java"],
+    tags = [
+    ],
+    deps = [
+        ":mutable_source",
+        "//common",
+        "//common/ast",
+        "//common/ast:mutable_expr",
+        "//common/types:type_providers",
+    ],
+)
+
+java_library(
+    name = "mutable_source",
+    srcs = ["CelMutableSource.java"],
+    tags = [
+    ],
+    deps = [
+        ":common",
+        "//:auto_value",
+        "//common/ast:mutable_expr",
+        "@maven//:com_google_errorprone_error_prone_annotations",
+        "@maven//:com_google_guava_guava",
+    ],
+)

common/src/main/java/dev/cel/common/ast/CelMutableAst.java → common/src/main/java/dev/cel/common/CelMutableAst.java

@@ -12,10 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package dev.cel.common.ast;
+package dev.cel.common;
 
-import dev.cel.common.CelAbstractSyntaxTree;
-import dev.cel.common.CelSource;
+import dev.cel.common.ast.CelMutableExpr;
+import dev.cel.common.ast.CelMutableExprConverter;
+import dev.cel.common.ast.CelReference;
 import dev.cel.common.types.CelType;
 import java.util.HashMap;
 import java.util.Map;
@@ -30,7 +31,7 @@
  */
 public final class CelMutableAst {
   private final CelMutableExpr mutatedExpr;
-  private final CelSource.Builder source;
+  private final CelMutableSource source;
   private final Map<Long, CelReference> references;
   private final Map<Long, CelType> types;
 
@@ -40,9 +41,10 @@ public CelMutableExpr expr() {
   }
 
   /**
-   * Returns the {@link CelSource} that was used during construction of the abstract syntax tree.
+   * Returns the {@link CelMutableSource} that was used during construction of the abstract syntax
+   * tree.
    */
-  public CelSource.Builder source() {
+  public CelMutableSource source() {
     return source;
   }
 
@@ -69,7 +71,7 @@ public Optional<CelType> getType(long exprId) {
   /** Converts this mutable AST into a parsed {@link CelAbstractSyntaxTree}. */
   public CelAbstractSyntaxTree toParsedAst() {
     return CelAbstractSyntaxTree.newParsedAst(
-        CelMutableExprConverter.fromMutableExpr(mutatedExpr), source.build());
+        CelMutableExprConverter.fromMutableExpr(mutatedExpr), source.toCelSource());
   }
 
   /**
@@ -79,7 +81,7 @@ public CelAbstractSyntaxTree toParsedAst() {
   public static CelMutableAst fromCelAst(CelAbstractSyntaxTree ast) {
     return new CelMutableAst(
         CelMutableExprConverter.fromCelExpr(ast.getExpr()),
-        ast.getSource().toBuilder(),
+        CelMutableSource.fromCelSource(ast.getSource()),
         ast.getReferenceMap(),
         ast.getTypeMap());
   }
@@ -88,21 +90,21 @@ public static CelMutableAst fromCelAst(CelAbstractSyntaxTree ast) {
    * 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);
+  public static CelMutableAst of(CelMutableExpr mutableExpr, CelMutableSource mutableSource) {
+    return new CelMutableAst(mutableExpr, mutableSource);
   }
 
-  private CelMutableAst(CelMutableExpr mutatedExpr, CelSource.Builder source) {
-    this(mutatedExpr, source, new HashMap<>(), new HashMap<>());
+  private CelMutableAst(CelMutableExpr mutatedExpr, CelMutableSource mutableSource) {
+    this(mutatedExpr, mutableSource, new HashMap<>(), new HashMap<>());
   }
 
   private CelMutableAst(
       CelMutableExpr mutatedExpr,
-      CelSource.Builder source,
+      CelMutableSource mutableSource,
       Map<Long, CelReference> references,
       Map<Long, CelType> types) {
     this.mutatedExpr = mutatedExpr;
-    this.source = source;
+    this.source = mutableSource;
     this.references = new HashMap<>(references);
     this.types = new HashMap<>(types);
   }

common/src/main/java/dev/cel/common/CelMutableSource.java

@@ -0,0 +1,114 @@
+// 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;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import dev.cel.common.CelSource.Extension;
+import dev.cel.common.ast.CelMutableExpr;
+import dev.cel.common.ast.CelMutableExprConverter;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Represents the mutable portion of the {@link CelSource}. This is intended for the purposes of
+ * augmenting an AST through CEL optimizers.
+ */
+public final class CelMutableSource {
+
+  final Map<Long, CelMutableExpr> macroCalls;
+  final Set<Extension> extensions;
+
+  @CanIgnoreReturnValue
+  public CelMutableSource addMacroCalls(long exprId, CelMutableExpr expr) {
+    this.macroCalls.put(exprId, checkNotNull(CelMutableExpr.newInstance(expr)));
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public CelMutableSource addAllMacroCalls(Map<Long, CelMutableExpr> macroCalls) {
+    this.macroCalls.putAll(macroCalls);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public CelMutableSource addAllExtensions(Collection<? extends Extension> extensions) {
+    checkNotNull(extensions);
+    this.extensions.addAll(extensions);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public CelMutableSource clearMacroCall(long exprId) {
+    this.macroCalls.remove(exprId);
+    return this;
+  }
+
+  @CanIgnoreReturnValue
+  public CelMutableSource clearMacroCalls() {
+    this.macroCalls.clear();
+    return this;
+  }
+
+  public Map<Long, CelMutableExpr> getMacroCalls() {
+    return macroCalls;
+  }
+
+  public Set<Extension> getExtensions() {
+    return extensions;
+  }
+
+  public CelSource toCelSource() {
+    return CelSource.newBuilder()
+        .addAllExtensions(extensions)
+        .addAllMacroCalls(
+            macroCalls.entrySet().stream()
+                .collect(
+                    toImmutableMap(
+                        Entry::getKey, v -> CelMutableExprConverter.fromMutableExpr(v.getValue()))))
+        .build();
+  }
+
+  public static CelMutableSource newInstance() {
+    return new CelMutableSource(new HashMap<>(), new HashSet<>());
+  }
+
+  public static CelMutableSource fromCelSource(CelSource source) {
+    return new CelMutableSource(
+        source.getMacroCalls().entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Entry::getKey,
+                    v -> CelMutableExprConverter.fromCelExpr(v.getValue()),
+                    (prev, next) -> {
+                      throw new IllegalStateException(
+                          "Unexpected source collision at ID: " + prev.id());
+                    },
+                    HashMap::new)),
+        source.getExtensions());
+  }
+
+  CelMutableSource(Map<Long, CelMutableExpr> macroCalls, Set<Extension> extensions) {
+    this.macroCalls = checkNotNull(macroCalls);
+    this.extensions = checkNotNull(extensions);
+  }
+}

common/src/main/java/dev/cel/common/CelSource.java

@@ -279,12 +279,6 @@ public Builder clearMacroCall(long exprId) {
       return this;
     }
 
-    @CanIgnoreReturnValue
-    public Builder clearMacroCalls() {
-      this.macroCalls.clear();
-      return this;
-    }
-
     public ImmutableSet<Extension> getExtensions() {
       return extensions.build();
     }

common/src/main/java/dev/cel/common/ast/BUILD.bazel

@@ -33,8 +33,7 @@ EXPR_FACTORY_SOURCES = [
 ]
 
 # keep sorted
-MUTABLE_AST_SOURCES = [
-    "CelMutableAst.java",
+MUTABLE_EXPR_SOURCES = [
     "CelMutableExpr.java",
     "CelMutableExprConverter.java",
 ]
@@ -119,14 +118,12 @@ java_library(
 )
 
 java_library(
-    name = "mutable_ast",
-    srcs = MUTABLE_AST_SOURCES,
+    name = "mutable_expr",
+    srcs = MUTABLE_EXPR_SOURCES,
     tags = [
     ],
     deps = [
         ":ast",
-        "//common",
-        "//common/types:type_providers",
         "@maven//:com_google_guava_guava",
     ],
 )

common/src/main/java/dev/cel/common/ast/CelMutableExpr.java

@@ -1064,6 +1064,9 @@ private CelMutableExpr(CelMutableExpr other) {
     this.id = other.id;
     this.exprKind = other.exprKind;
     switch (other.getKind()) {
+      case NOT_SET:
+        this.exprValue = CelExpr.newBuilder().build().exprKind().notSet();
+        break;
       case CONSTANT:
         this.exprValue = other.exprValue; // Constant is immutable.
         break;

common/src/main/java/dev/cel/common/ast/CelMutableExprConverter.java

@@ -134,8 +134,6 @@ private static CelMutableCreateMap fromCelMapToMutableMap(CelCreateMap celCreate
     return CelMutableCreateMap.create(entries);
   }
 
-  ///////////////////////
-
   public static CelExpr fromMutableExpr(CelMutableExpr mutableExpr) {
     long id = mutableExpr.id();
     switch (mutableExpr.getKind()) {

common/src/main/java/dev/cel/common/navigation/BUILD.bazel

@@ -56,7 +56,8 @@ java_library(
     deps = [
         ":common",
         "//:auto_value",
-        "//common/ast:mutable_ast",
+        "//common:mutable_ast",
+        "//common/ast:mutable_expr",
         "//common/types:type_providers",
         "@maven//:com_google_errorprone_error_prone_annotations",
         "@maven//:com_google_guava_guava",

common/src/main/java/dev/cel/common/navigation/CelNavigableMutableAst.java

@@ -14,7 +14,7 @@
 
 package dev.cel.common.navigation;
 
-import dev.cel.common.ast.CelMutableAst;
+import dev.cel.common.CelMutableAst;
 import dev.cel.common.types.CelType;
 import java.util.Optional;
 

common/src/test/java/dev/cel/common/ast/BUILD.bazel

@@ -13,13 +13,15 @@ java_library(
         "//:java_truth",
         "//common",
         "//common:compiler_common",
+        "//common:mutable_ast",
+        "//common:mutable_source",
         "//common:options",
         "//common/ast",
         "//common/ast:cel_expr_visitor",
         "//common/ast:expr_converter",
         "//common/ast:expr_factory",
         "//common/ast:expr_v1alpha1_converter",
-        "//common/ast:mutable_ast",
+        "//common/ast:mutable_expr",
         "//common/resources/testdata/proto3:test_all_types_java_proto",
         "//common/types",
         "//compiler",

common/src/test/java/dev/cel/common/ast/CelMutableAstTest.java

@@ -17,8 +17,9 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import dev.cel.common.CelAbstractSyntaxTree;
+import dev.cel.common.CelMutableAst;
+import dev.cel.common.CelMutableSource;
 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;
@@ -33,12 +34,12 @@ public final class CelMutableAstTest {
   @Test
   public void constructMutableAst() {
     CelMutableExpr mutableExpr = CelMutableExpr.ofConstant(1L, CelConstant.ofValue("hello world"));
-    CelSource.Builder sourceBuilder = CelSource.newBuilder();
+    CelMutableSource mutableSource = CelMutableSource.newInstance();
 
-    CelMutableAst celMutableAst = CelMutableAst.of(mutableExpr, sourceBuilder);
+    CelMutableAst celMutableAst = CelMutableAst.of(mutableExpr, mutableSource);
 
     assertThat(celMutableAst.expr()).isEqualTo(mutableExpr);
-    assertThat(celMutableAst.source()).isSameInstanceAs(sourceBuilder);
+    assertThat(celMutableAst.source()).isSameInstanceAs(mutableSource);
   }
 
   @Test

common/src/test/java/dev/cel/common/navigation/BUILD.bazel

@@ -10,9 +10,10 @@ java_library(
         "//:java_truth",
         "//common",
         "//common:compiler_common",
+        "//common:mutable_ast",
         "//common:options",
         "//common/ast",
-        "//common/ast:mutable_ast",
+        "//common/ast:mutable_expr",
         "//common/navigation",
         "//common/navigation:common",
         "//common/navigation:mutable_navigation",

common/src/test/java/dev/cel/common/navigation/CelNavigableMutableAstTest.java

@@ -16,8 +16,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import dev.cel.common.CelMutableAst;
 import dev.cel.common.ast.CelConstant;
-import dev.cel.common.ast.CelMutableAst;
 import dev.cel.common.ast.CelMutableExpr;
 import dev.cel.compiler.CelCompiler;
 import dev.cel.compiler.CelCompilerFactory;