Allow MutableAst to insert a new bind macro

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

@@ -88,6 +88,7 @@ java_library(
     ],
     deps = [
         ":ast",
+        "//common/annotations",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java",
     ],

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

@@ -216,6 +216,7 @@ public CelComprehension comprehensionOrDefault() {
   /** Builder for CelExpr. */
   @AutoValue.Builder
   public abstract static class Builder {
+
     public abstract long id();
 
     public abstract Builder setId(long value);
@@ -787,6 +788,8 @@ public abstract static class Entry {
       @AutoValue.Builder
       public abstract static class Builder {
 
+        public abstract long id();
+
         public abstract CelExpr value();
 
         public abstract Builder setId(long value);
@@ -918,6 +921,7 @@ public abstract static class Entry {
       /** Builder for CelCreateMap.Entry. */
       @AutoValue.Builder
       public abstract static class Builder {
+        public abstract long id();
 
         public abstract CelExpr key();
 

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

@@ -19,16 +19,24 @@
 
 import com.google.common.primitives.UnsignedLong;
 import com.google.protobuf.ByteString;
+import dev.cel.common.annotations.Internal;
 import java.util.Arrays;
 
 /** Factory for generating expression nodes. */
+@Internal
 public class CelExprFactory {
-  private final CelExprIdGeneratorFactory.MonotonicIdGenerator idGenerator;
+
+  private final CelExprIdGeneratorFactory.ExprIdGenerator idGenerator;
 
   public static CelExprFactory newInstance() {
     return new CelExprFactory();
   }
 
+  public static CelExprFactory newInstance(
+      CelExprIdGeneratorFactory.ExprIdGenerator exprIdGenerator) {
+    return new CelExprFactory(exprIdGenerator);
+  }
+
   /** Create a new constant expression. */
   public final CelExpr newConstant(CelConstant constant) {
     return CelExpr.newBuilder().setId(nextExprId()).setConstant(constant).build();
@@ -543,10 +551,19 @@ public final CelExpr newSelect(CelExpr operand, String field, boolean testOnly)
 
   /** Returns the next unique expression ID. */
   protected long nextExprId() {
-    return idGenerator.nextExprId();
+    return idGenerator.generate(
+        /* exprId= */ -1); // Unconditionally generate next unique ID (i.e: no renumbering).
   }
 
   protected CelExprFactory() {
-    idGenerator = CelExprIdGeneratorFactory.newMonotonicIdGenerator(0);
+    this(CelExprIdGeneratorFactory.newMonotonicIdGenerator(0));
+  }
+
+  private CelExprFactory(CelExprIdGeneratorFactory.MonotonicIdGenerator idGenerator) {
+    this((unused) -> idGenerator.nextExprId());
+  }
+
+  private CelExprFactory(CelExprIdGeneratorFactory.ExprIdGenerator exprIdGenerator) {
+    idGenerator = exprIdGenerator;
   }
 }

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

@@ -14,11 +14,17 @@
 
 package dev.cel.common.ast;
 
+import com.google.common.collect.ImmutableSet;
+
 /** Provides string formatting support for {@link CelExpr}. */
 final class CelExprFormatter {
   private final StringBuilder indent = new StringBuilder();
   private final StringBuilder exprBuilder = new StringBuilder();
 
+  /** Denotes a set of expression kinds that will not have a new line inserted. */
+  private static final ImmutableSet<CelExpr.ExprKind.Kind> EXCLUDED_NEWLINE_KINDS =
+      ImmutableSet.of(CelExpr.ExprKind.Kind.CONSTANT, CelExpr.ExprKind.Kind.NOT_SET);
+
   static String format(CelExpr celExpr) {
     CelExprFormatter formatter = new CelExprFormatter();
     formatter.formatExpr(celExpr);
@@ -28,7 +34,7 @@ static String format(CelExpr celExpr) {
   private void formatExpr(CelExpr celExpr) {
     append(String.format("%s [%d] {", celExpr.exprKind().getKind(), celExpr.id()));
     CelExpr.ExprKind.Kind exprKind = celExpr.exprKind().getKind();
-    if (!exprKind.equals(CelExpr.ExprKind.Kind.CONSTANT)) {
+    if (!EXCLUDED_NEWLINE_KINDS.contains(exprKind)) {
       appendNewline();
     }
 
@@ -57,12 +63,17 @@ private void formatExpr(CelExpr celExpr) {
       case COMPREHENSION:
         appendComprehension(celExpr.comprehension());
         break;
+      case NOT_SET:
+        break;
       default:
+        // This should be unreachable unless if we've added any other kinds.
+        indent();
         append("Unknown kind: " + exprKind);
+        outdent();
         break;
     }
 
-    if (!exprKind.equals(CelExpr.ExprKind.Kind.CONSTANT)) {
+    if (!EXCLUDED_NEWLINE_KINDS.contains(exprKind)) {
       appendNewline();
       append("}");
     } else {

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

@@ -59,8 +59,45 @@ public static class StableIdGenerator {
     private final HashMap<Long, Long> idSet;
     private long exprId;
 
+    /** Checks if the given ID has been encountered before. */
+    public boolean hasId(long id) {
+      return idSet.containsKey(id);
+    }
+
+    /**
+     * Generate the next available ID while memoizing the existing ID.
+     *
+     * <p>The main purpose of this is to sanitize a new AST to replace an existing AST's node with.
+     * The incoming AST may not have its IDs consistently numbered (often, the expr IDs are just
+     * zeroes). In those cases, we just want to return an incremented expr ID.
+     *
+     * <p>The memoization becomes necessary if the incoming AST contains an expression with macro
+     * map populated, requiring a normalization pass. In this case, the method behaves largely the
+     * same as {@link #renumberId}.
+     *
+     * @param id Existing ID to memoize. Providing 0 or less will skip the memoization, in which
+     *     case this behaves just like a {@link MonotonicIdGenerator}.
+     */
+    public long nextExprId(long id) {
+      long nextExprId = ++exprId;
+      if (id > 0) {
+        idSet.put(id, nextExprId);
+      }
+      return nextExprId;
+    }
+
+    /** Memoize a given expression ID with a newly generated ID. */
+    public void memoize(long existingId, long newId) {
+      idSet.put(existingId, newId);
+    }
+
+    /**
+     * Renumbers the existing expression ID to a newly generated unique ID. The existing ID is
+     * memoized, and calling this method again with the same ID will always return the same
+     * generated ID.
+     */
     public long renumberId(long id) {
-      Preconditions.checkArgument(id >= 0);
+      Preconditions.checkArgument(id >= 0, "Expr ID must be positive. Got: %s", id);
       if (id == 0) {
         return 0;
       }
@@ -81,5 +118,13 @@ private StableIdGenerator(long exprId) {
     }
   }
 
+  /** Functional interface for generating the next unique expression ID. */
+  @FunctionalInterface
+  public interface ExprIdGenerator {
+
+    /** Generates an expression ID with the provided expr ID as the context. */
+    long generate(long exprId);
+  }
+
   private CelExprIdGeneratorFactory() {}
 }

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

@@ -16,6 +16,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import dev.cel.common.ast.CelExprIdGeneratorFactory.StableIdGenerator;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -37,4 +38,15 @@ public void nextExprId_startingDefaultIsOne() {
     assertThat(exprFactory.nextExprId()).isEqualTo(1L);
     assertThat(exprFactory.nextExprId()).isEqualTo(2L);
   }
+
+  @Test
+  public void nextExprId_usingStableIdGenerator() {
+    StableIdGenerator stableIdGenerator = CelExprIdGeneratorFactory.newStableIdGenerator(0);
+    CelExprFactory exprFactory = CelExprFactory.newInstance(stableIdGenerator::nextExprId);
+
+    assertThat(exprFactory.nextExprId()).isEqualTo(1L);
+    assertThat(exprFactory.nextExprId()).isEqualTo(2L);
+    assertThat(stableIdGenerator.hasId(-1)).isFalse();
+    assertThat(stableIdGenerator.hasId(0)).isFalse();
+  }
 }

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

@@ -61,6 +61,13 @@ public void constant(@TestParameter ConstantTestCase constantTestCase) throws Ex
     assertThat(formattedExpr).isEqualTo(constantTestCase.formatted);
   }
 
+  @Test
+  public void notSet() {
+    String formattedExpr = CelExprFormatter.format(CelExpr.ofNotSet(1));
+
+    assertThat(formattedExpr).isEqualTo("NOT_SET [1] {}");
+  }
+
   @Test
   public void select() throws Exception {
     CelCompiler celCompiler =

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

@@ -81,6 +81,7 @@ java_library(
     tags = [
     ],
     deps = [
+        "//:auto_value",
         "//common",
         "//common/annotations",
         "//common/ast",