Use mutable expressions for AST rewrites during type-check

Author: l46kok

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

@@ -179,12 +179,14 @@ java_library(
         "//common:cel_ast",
         "//common:compiler_common",
         "//common:container",
+        "//common:mutable_ast",
         "//common:operator",
         "//common:options",
         "//common:proto_ast",
         "//common/annotations",
         "//common/ast",
         "//common/ast:expr_converter",
+        "//common/ast:mutable_expr",
         "//common/internal:errors",
         "//common/internal:file_descriptor_converter",
         "//common/types",

checker/src/main/java/dev/cel/checker/Env.java

@@ -38,6 +38,7 @@
 import dev.cel.common.ast.CelConstant;
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.ast.CelExprConverter;
+import dev.cel.common.ast.CelMutableExpr;
 import dev.cel.common.ast.CelReference;
 import dev.cel.common.internal.Errors;
 import dev.cel.common.types.CelKind;
@@ -304,6 +305,14 @@ public CelType getType(CelExpr expr) {
     return Preconditions.checkNotNull(typeMap.get(expr.id()), "expression has no type");
   }
 
+  /**
+   * Returns the type associated with a mutable expression by expression id. It's an error to call this
+   * method if the type is not present.
+   */
+  CelType getType(CelMutableExpr expr) {
+    return Preconditions.checkNotNull(typeMap.get(expr.id()), "expression has no type");
+  }
+
   /**
    * Sets the type associated with an expression by id. It's an error if the type is already set and
    * is different than the provided one. Returns the expression parameter.
@@ -319,6 +328,21 @@ public CelExpr setType(CelExpr expr, CelType type) {
     return expr;
   }
 
+  /**
+   * Sets the type associated with a mutable expression by id. It's an error if the type is already set and
+   * is different than the provided one. Returns the expression parameter.
+   */
+  @CanIgnoreReturnValue
+  CelMutableExpr setType(CelMutableExpr expr, CelType type) {
+    CelType oldType = typeMap.put(expr.id(), type);
+    Preconditions.checkState(
+        oldType == null || oldType.equals(type),
+        "expression already has a type which is incompatible.\n old: %s\n new: %s",
+        oldType,
+        type);
+    return expr;
+  }
+
   /**
    * Sets the reference associated with an expression. It's an error if the reference is already set
    * and is different.
@@ -330,6 +354,17 @@ public void setRef(CelExpr expr, CelReference reference) {
         "expression already has a reference which is incompatible");
   }
 
+  /**
+   * Sets the reference associated with a mutable expression. It's an error if the reference is already set
+   * and is different.
+   */
+  void setRef(CelMutableExpr expr, CelReference reference) {
+    CelReference oldReference = referenceMap.put(expr.id(), reference);
+    Preconditions.checkState(
+        oldReference == null || oldReference.equals(reference),
+        "expression already has a reference which is incompatible");
+  }
+
   /**
    * Adds a declaration to the environment, based on the Decl proto. Will report errors if the
    * declaration overlaps with an existing one, or clashes with a macro.