From cb53abe25b77771100909254b5f76a769fe5c807 Mon Sep 17 00:00:00 2001 From: Sean Huh Date: Tue, 13 Oct 2026 10:27:33 -0700 Subject: [PATCH] Cache ExprPropertyCalculator on BaseNavigableExpr and key properties by identity PiperOrigin-RevId: 998727839 --- .../common/navigation/BaseNavigableExpr.java | 23 ++++++++ .../common/navigation/CelNavigableExpr.java | 10 +--- .../navigation/CelNavigableExprVisitor.java | 54 +++++++++---------- .../navigation/CelNavigableMutableExpr.java | 11 +--- .../navigation/ExprPropertyCalculator.java | 34 +++++------- .../CelNavigableMutableExprTest.java | 21 ++++++++ 6 files changed, 84 insertions(+), 69 deletions(-) diff --git a/common/src/main/java/dev/cel/common/navigation/BaseNavigableExpr.java b/common/src/main/java/dev/cel/common/navigation/BaseNavigableExpr.java index dabcac3a2..8eff19fc5 100644 --- a/common/src/main/java/dev/cel/common/navigation/BaseNavigableExpr.java +++ b/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 { + private ExprPropertyCalculator exprPropertyCalculator; + BaseNavigableExpr() {} + ExprPropertyCalculator exprPropertyCalculator() { + if (exprPropertyCalculator == null) { + exprPropertyCalculator = new ExprPropertyCalculator<>(expr()); + } + return exprPropertyCalculator; + } + + static > T buildWithProperties( + Builder builder, E expr, ExprPropertyCalculator exprPropertyCalculator) { + ExprProperty exprProperty = exprPropertyCalculator.getProperty(expr); + T navigableExpr = + builder + .setExpr(expr) + .setHeight(exprProperty.height()) + .setMaxId(exprProperty.maxId()) + .build(); + ((BaseNavigableExpr) navigableExpr).exprPropertyCalculator = exprPropertyCalculator; + return navigableExpr; + } + public abstract E expr(); public long id() { diff --git a/common/src/main/java/dev/cel/common/navigation/CelNavigableExpr.java b/common/src/main/java/dev/cel/common/navigation/CelNavigableExpr.java index 69453b900..623855fa6 100644 --- a/common/src/main/java/dev/cel/common/navigation/CelNavigableExpr.java +++ b/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 { /** Constructs a new instance of {@link CelNavigableExpr} from {@link CelExpr}. */ public static CelNavigableExpr fromExpr(CelExpr expr) { - ExprPropertyCalculator 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 diff --git a/common/src/main/java/dev/cel/common/navigation/CelNavigableExprVisitor.java b/common/src/main/java/dev/cel/common/navigation/CelNavigableExprVisitor.java index c8eb57d53..66dc6cdee 100644 --- a/common/src/main/java/dev/cel/common/navigation/CelNavigableExprVisitor.java +++ b/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 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,10 +72,9 @@ static > Stream collect( */ static > Stream collect( T navigableExpr, int maxDepth, TraversalOrder traversalOrder) { - ExprPropertyCalculator exprHeightCalculator = - new ExprPropertyCalculator<>(navigableExpr.expr()); CelNavigableExprVisitor visitor = - new CelNavigableExprVisitor<>(maxDepth, exprHeightCalculator, traversalOrder); + new CelNavigableExprVisitor<>( + maxDepth, navigableExpr.exprPropertyCalculator(), traversalOrder); visitor.visit(navigableExpr); @@ -92,12 +82,14 @@ static > Stream collect( } 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 list, T parent) { } private T newNavigableChild(T parent, E expr) { - ExprProperty exprProperty = exprPropertyCalculator.getProperty(expr.id()); - - return parent - .builderFromInstance() - .setExpr(expr) - .setDepth(parent.depth() + 1) - .setHeight(exprProperty.height()) - .setMaxId(exprProperty.maxId()) - .setParent(parent) - .build(); + return BaseNavigableExpr.buildWithProperties( + parent.builderFromInstance().setDepth(parent.depth() + 1).setParent(parent), + expr, + exprPropertyCalculator); + } + + private CelNavigableExprVisitor( + int maxDepth, + ExprPropertyCalculator exprPropertyCalculator, + TraversalOrder traversalOrder) { + this.maxDepth = maxDepth; + this.exprPropertyCalculator = exprPropertyCalculator; + this.traversalOrder = traversalOrder; + this.streamBuilder = Stream.builder(); } } diff --git a/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableExpr.java b/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableExpr.java index b9af67675..fc7b362d3 100644 --- a/common/src/main/java/dev/cel/common/navigation/CelNavigableMutableExpr.java +++ b/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 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 diff --git a/common/src/main/java/dev/cel/common/navigation/ExprPropertyCalculator.java b/common/src/main/java/dev/cel/common/navigation/ExprPropertyCalculator.java index be2c07221..5f6eeab38 100644 --- a/common/src/main/java/dev/cel/common/navigation/ExprPropertyCalculator.java +++ b/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 { - // Store hashmap instead of immutable map for performance, such that this helper class can be - // instantiated faster. - private final HashMap 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 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 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 { diff --git a/common/src/test/java/dev/cel/common/navigation/CelNavigableMutableExprTest.java b/common/src/test/java/dev/cel/common/navigation/CelNavigableMutableExprTest.java index 6900d41aa..7a6acf3bd 100644 --- a/common/src/test/java/dev/cel/common/navigation/CelNavigableMutableExprTest.java +++ b/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 heights = + mutableExpr + .allNodes(TraversalOrder.PRE_ORDER) + .map(BaseNavigableExpr::height) + .collect(toImmutableList()); + + assertThat(heights).containsExactly(2, 0, 1, 0).inOrder(); + } }