Compute Max IDs per subtree in NavigableExpr

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

@@ -14,7 +14,7 @@ java_library(
         "CelNavigableAst.java",
         "CelNavigableExpr.java",
         "CelNavigableExprVisitor.java",
-        "ExprHeightCalculator.java",
+        "ExprPropertyCalculator.java",
         "TraversalOrder.java",
     ],
     tags = [

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

@@ -40,6 +40,12 @@ public long id() {
   /** Represents the count of transitive parents. Depth of an AST's root is 0. */
   public abstract int depth();
 
+  /**
+   * Represents the maximum ID of the tree. Note that if the underlying expression tree held by this
+   * navigable expression is mutated, its max ID becomes stale and must be recomputed.
+   */
+  public abstract long maxId();
+
   /**
    * Represents the maximum count of children from any of its branches. Height of a leaf node is 0.
    * For example, the height of the call node 'func' in expression `(1 + 2 + 3).func(4 + 5)` is 3.
@@ -125,6 +131,9 @@ default ExprKind.Kind getKind() {
     @CanIgnoreReturnValue
     Builder<E, T> setHeight(int value);
 
+    @CanIgnoreReturnValue
+    Builder<E, T> setMaxId(long value);
+
     @CheckReturnValue
     T build();
   }

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

@@ -16,6 +16,7 @@
 
 import com.google.auto.value.AutoValue;
 import dev.cel.common.ast.CelExpr;
+import dev.cel.common.navigation.ExprPropertyCalculator.ExprProperty;
 import java.util.Optional;
 import java.util.stream.Stream;
 
@@ -28,10 +29,17 @@
 // redundant override: Overriding is required to specify the return type to a concrete type.
 @SuppressWarnings({"unchecked", "RedundantOverride"})
 public abstract class CelNavigableExpr extends BaseNavigableExpr<CelExpr> {
+
   /** Constructs a new instance of {@link CelNavigableExpr} from {@link CelExpr}. */
   public static CelNavigableExpr fromExpr(CelExpr expr) {
-    ExprHeightCalculator<CelExpr> exprHeightCalculator = new ExprHeightCalculator<>(expr);
-    return builder().setExpr(expr).setHeight(exprHeightCalculator.getHeight(expr.id())).build();
+    ExprPropertyCalculator<CelExpr> exprHeightCalculator = new ExprPropertyCalculator<>(expr);
+    ExprProperty exprProperty = exprHeightCalculator.getProperty(expr.id());
+
+    return builder()
+        .setExpr(expr)
+        .setHeight(exprProperty.height())
+        .setMaxId(exprProperty.maxId())
+        .build();
   }
 
   @Override
@@ -74,7 +82,7 @@ public Builder builderFromInstance() {
 
   /** Create a new builder to construct a {@link CelNavigableExpr} instance. */
   public static Builder builder() {
-    return new AutoValue_CelNavigableExpr.Builder().setDepth(0).setHeight(0);
+    return new AutoValue_CelNavigableExpr.Builder().setDepth(0).setHeight(0).setMaxId(0);
   }
 
   /** Builder to configure {@link CelNavigableExpr}. */
@@ -91,6 +99,9 @@ public abstract static class Builder
     @Override
     public abstract Builder setDepth(int value);
 
+    @Override
+    public abstract Builder setMaxId(long value);
+
     @Override
     public abstract Builder setHeight(int value);
   }

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

@@ -15,6 +15,7 @@
 package dev.cel.common.navigation;
 
 import dev.cel.common.ast.Expression;
+import dev.cel.common.navigation.ExprPropertyCalculator.ExprProperty;
 import java.util.List;
 import java.util.stream.Stream;
 
@@ -23,12 +24,14 @@ final class CelNavigableExprVisitor<E extends Expression, T extends BaseNavigabl
   private static final int MAX_DESCENDANTS_RECURSION_DEPTH = 500;
 
   private final Stream.Builder<T> streamBuilder;
-  private final ExprHeightCalculator<E> exprPropertyCalculator;
+  private final ExprPropertyCalculator<E> exprPropertyCalculator;
   private final TraversalOrder traversalOrder;
   private final int maxDepth;
 
   private CelNavigableExprVisitor(
-      int maxDepth, ExprHeightCalculator<E> exprPropertyCalculator, TraversalOrder traversalOrder) {
+      int maxDepth,
+      ExprPropertyCalculator<E> exprPropertyCalculator,
+      TraversalOrder traversalOrder) {
     this.maxDepth = maxDepth;
     this.exprPropertyCalculator = exprPropertyCalculator;
     this.traversalOrder = traversalOrder;
@@ -78,7 +81,8 @@ static <E extends Expression, T extends BaseNavigableExpr<E>> Stream<T> collect(
    */
   static <E extends Expression, T extends BaseNavigableExpr<E>> Stream<T> collect(
       T navigableExpr, int maxDepth, TraversalOrder traversalOrder) {
-    ExprHeightCalculator<E> exprHeightCalculator = new ExprHeightCalculator<>(navigableExpr.expr());
+    ExprPropertyCalculator<E> exprHeightCalculator =
+        new ExprPropertyCalculator<>(navigableExpr.expr());
     CelNavigableExprVisitor<E, T> visitor =
         new CelNavigableExprVisitor<>(maxDepth, exprHeightCalculator, traversalOrder);
 
@@ -174,11 +178,14 @@ private void visitExprList(List<E> createListExpr, T parent) {
   }
 
   private T newNavigableChild(T parent, E expr) {
+    ExprProperty exprProperty = exprPropertyCalculator.getProperty(expr.id());
+
     return parent
         .<T>builderFromInstance()
         .setExpr(expr)
         .setDepth(parent.depth() + 1)
-        .setHeight(exprPropertyCalculator.getHeight(expr.id()))
+        .setHeight(exprProperty.height())
+        .setMaxId(exprProperty.maxId())
         .setParent(parent)
         .build();
   }