Cache ExprPropertyCalculator on BaseNavigableExpr and key properties by identity

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

@@ -20,6 +20,7 @@
 import dev.cel.common.ast.CelExpr;
 import dev.cel.common.ast.CelExpr.ExprKind;
 import dev.cel.common.ast.Expression;
+import dev.cel.common.navigation.ExprPropertyCalculator.ExprProperty;
 import java.util.Optional;
 import java.util.stream.Stream;
 
@@ -34,8 +35,30 @@
 @SuppressWarnings("unchecked") // Generic types are properly bound to Expression
 public abstract class BaseNavigableExpr<E extends Expression> {
 
+  private ExprPropertyCalculator<E> exprPropertyCalculator;
+
   BaseNavigableExpr() {}
 
+  ExprPropertyCalculator<E> exprPropertyCalculator() {
+    if (exprPropertyCalculator == null) {
+      exprPropertyCalculator = new ExprPropertyCalculator<>(expr());
+    }
+    return exprPropertyCalculator;
+  }
+
+  static <E extends Expression, T extends BaseNavigableExpr<E>> T buildWithProperties(
+      Builder<E, T> builder, E expr, ExprPropertyCalculator<E> exprPropertyCalculator) {
+    ExprProperty exprProperty = exprPropertyCalculator.getProperty(expr);
+    T navigableExpr =
+        builder
+            .setExpr(expr)
+            .setHeight(exprProperty.height())
+            .setMaxId(exprProperty.maxId())
+            .build();
+    ((BaseNavigableExpr<E>) navigableExpr).exprPropertyCalculator = exprPropertyCalculator;
+    return navigableExpr;
+  }
+
   public abstract E expr();
 
   public long id() {

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

@@ -16,7 +16,6 @@
 
 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;
 
@@ -32,14 +31,7 @@ public abstract class CelNavigableExpr extends BaseNavigableExpr<CelExpr> {
 
   /** Constructs a new instance of {@link CelNavigableExpr} from {@link CelExpr}. */
   public static CelNavigableExpr fromExpr(CelExpr expr) {
-    ExprPropertyCalculator<CelExpr> exprHeightCalculator = new ExprPropertyCalculator<>(expr);
-    ExprProperty exprProperty = exprHeightCalculator.getProperty(expr.id());
-
-    return builder()
-        .setExpr(expr)
-        .setHeight(exprProperty.height())
-        .setMaxId(exprProperty.maxId())
-        .build();
+    return buildWithProperties(builder(), expr, new ExprPropertyCalculator<>(expr));
   }
 
   @Override

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

@@ -14,9 +14,10 @@
 
 package dev.cel.common.navigation;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import dev.cel.common.ast.Expression;
 import dev.cel.common.ast.Expression.List;
-import dev.cel.common.navigation.ExprPropertyCalculator.ExprProperty;
 import java.util.stream.Stream;
 
 /** Visitor implementation to navigate an AST. */
@@ -28,16 +29,6 @@ final class CelNavigableExprVisitor<E extends Expression, T extends BaseNavigabl
   private final TraversalOrder traversalOrder;
   private final int maxDepth;
 
-  private CelNavigableExprVisitor(
-      int maxDepth,
-      ExprPropertyCalculator<E> exprPropertyCalculator,
-      TraversalOrder traversalOrder) {
-    this.maxDepth = maxDepth;
-    this.exprPropertyCalculator = exprPropertyCalculator;
-    this.traversalOrder = traversalOrder;
-    this.streamBuilder = Stream.builder();
-  }
-
   /**
    * Returns a stream containing all the nodes using the specified traversal order. Note that the
    * collection occurs eagerly.
@@ -81,23 +72,24 @@ 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) {
-    ExprPropertyCalculator<E> exprHeightCalculator =
-        new ExprPropertyCalculator<>(navigableExpr.expr());
     CelNavigableExprVisitor<E, T> visitor =
-        new CelNavigableExprVisitor<>(maxDepth, exprHeightCalculator, traversalOrder);
+        new CelNavigableExprVisitor<>(
+            maxDepth, navigableExpr.exprPropertyCalculator(), traversalOrder);
 
     visitor.visit(navigableExpr);
 
     return visitor.streamBuilder.build();
   }
 
   private void visit(T navigableExpr) {
-    if (navigableExpr.depth() > MAX_DESCENDANTS_RECURSION_DEPTH - 1) {
-      throw new IllegalStateException("Max recursion depth reached.");
+    checkState(
+        navigableExpr.depth() <= MAX_DESCENDANTS_RECURSION_DEPTH - 1,
+        "Max recursion depth reached.");
+    if (navigableExpr.depth() > maxDepth) {
+      return;
     }
 
-    boolean addToStream = navigableExpr.depth() <= maxDepth;
-    if (addToStream && traversalOrder.equals(TraversalOrder.PRE_ORDER)) {
+    if (traversalOrder.equals(TraversalOrder.PRE_ORDER)) {
       streamBuilder.add(navigableExpr);
     }
 
@@ -124,7 +116,7 @@ private void visit(T navigableExpr) {
         break;
     }
 
-    if (addToStream && traversalOrder.equals(TraversalOrder.POST_ORDER)) {
+    if (traversalOrder.equals(TraversalOrder.POST_ORDER)) {
       streamBuilder.add(navigableExpr);
     }
   }
@@ -177,15 +169,19 @@ private void visitExprList(java.util.List<E> list, 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(exprProperty.height())
-        .setMaxId(exprProperty.maxId())
-        .setParent(parent)
-        .build();
+    return BaseNavigableExpr.buildWithProperties(
+        parent.<T>builderFromInstance().setDepth(parent.depth() + 1).setParent(parent),
+        expr,
+        exprPropertyCalculator);
+  }
+
+  private CelNavigableExprVisitor(
+      int maxDepth,
+      ExprPropertyCalculator<E> exprPropertyCalculator,
+      TraversalOrder traversalOrder) {
+    this.maxDepth = maxDepth;
+    this.exprPropertyCalculator = exprPropertyCalculator;
+    this.traversalOrder = traversalOrder;
+    this.streamBuilder = Stream.builder();
   }
 }

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

@@ -16,7 +16,6 @@
 
 import com.google.auto.value.AutoValue;
 import dev.cel.common.ast.CelMutableExpr;
-import dev.cel.common.navigation.ExprPropertyCalculator.ExprProperty;
 import java.util.Optional;
 import java.util.stream.Stream;
 
@@ -32,15 +31,7 @@ public abstract class CelNavigableMutableExpr extends BaseNavigableExpr<CelMutab
 
   /** Constructs a new instance of {@link CelNavigableMutableExpr} from {@link CelMutableExpr}. */
   public static CelNavigableMutableExpr fromExpr(CelMutableExpr expr) {
-    ExprPropertyCalculator<CelMutableExpr> exprHeightCalculator =
-        new ExprPropertyCalculator<>(expr);
-    ExprProperty exprProperty = exprHeightCalculator.getProperty(expr.id());
-
-    return builder()
-        .setExpr(expr)
-        .setHeight(exprProperty.height())
-        .setMaxId(exprProperty.maxId())
-        .build();
+    return buildWithProperties(builder(), expr, new ExprPropertyCalculator<>(expr));
   }
 
   @Override

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

@@ -21,30 +21,17 @@
 import dev.cel.common.ast.Expression.List;
 import dev.cel.common.ast.Expression.Map;
 import dev.cel.common.ast.Expression.Struct;
-import java.util.HashMap;
+import java.util.IdentityHashMap;
 
 /** Package-private class to assist computing the height and the max ID of expression nodes. */
 final class ExprPropertyCalculator<E extends Expression> {
-  // Store hashmap instead of immutable map for performance, such that this helper class can be
-  // instantiated faster.
-  private final HashMap<Long, ExprProperty> idToProperty;
+  // Store IdentityHashMap instead of immutable map for performance, such that this helper class can
+  // be instantiated faster and nodes with duplicate IDs (e.g. 0) do not collide.
+  private final IdentityHashMap<E, ExprProperty> exprToProperty;
 
-  ExprPropertyCalculator(E celExpr) {
-    this.idToProperty = new HashMap<>();
-    visit(celExpr);
-  }
-
-  /**
-   * Retrieves the property containing the expression's maximum ID and the height of the subtree.
-   *
-   * @throws IllegalArgumentException If the provided expression ID does not exist.
-   */
-  ExprProperty getProperty(Long exprId) {
-    if (!idToProperty.containsKey(exprId)) {
-      throw new IllegalArgumentException("Property not found for expression id: " + exprId);
-    }
-
-    return idToProperty.get(exprId);
+  ExprProperty getProperty(E expr) {
+    ExprProperty property = exprToProperty.get(expr);
+    return property != null ? property : visit(expr);
   }
 
   private ExprProperty visit(E expr) {
@@ -79,7 +66,7 @@ private ExprProperty visit(E expr) {
     ExprProperty exprProperty =
         ExprProperty.create(
             baseHeight + visitedProperty.height(), max(visitedProperty.maxId(), expr.id()));
-    idToProperty.put(expr.id(), exprProperty);
+    exprToProperty.put(expr, exprProperty);
 
     return exprProperty;
   }
@@ -137,6 +124,11 @@ private ExprProperty visitExprList(java.util.List<E> list) {
     return visitedProperty;
   }
 
+  ExprPropertyCalculator(E celExpr) {
+    this.exprToProperty = new IdentityHashMap<>();
+    visit(celExpr);
+  }
+
   /** Value class to store the height and the max ID at a specific expression ID. */
   @AutoValue
   abstract static class ExprProperty {

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

@@ -39,6 +39,7 @@ public void construct_withoutParent() {
     assertThat(navigableExpr.expr()).isEqualTo(constExpr);
     assertThat(navigableExpr.depth()).isEqualTo(2);
     assertThat(navigableExpr.parent()).isEmpty();
+    assertThat(navigableExpr.allNodes()).containsExactly(navigableExpr);
   }
 
   @Test
@@ -147,4 +148,24 @@ public void children_filteredConstants_returnsSingleConstant() {
     assertThat(allNodes)
         .containsExactly(CelMutableExpr.ofConstant(CelConstant.ofValue("element1")));
   }
+
+  @Test
+  public void allNodes_zeroIds_heightsPopulated() {
+    CelNavigableMutableExpr mutableExpr =
+        CelNavigableMutableExpr.fromExpr(
+            CelMutableExpr.ofList(
+                CelMutableList.create(
+                    CelMutableExpr.ofConstant(CelConstant.ofValue("element1")),
+                    CelMutableExpr.ofList(
+                        CelMutableList.create(
+                            CelMutableExpr.ofConstant(CelConstant.ofValue("element2")))))));
+
+    ImmutableList<Integer> heights =
+        mutableExpr
+            .allNodes(TraversalOrder.PRE_ORDER)
+            .map(BaseNavigableExpr::height)
+            .collect(toImmutableList());
+
+    assertThat(heights).containsExactly(2, 0, 1, 0).inOrder();
+  }
 }