Add accessors for subtree size and height for navigable AST.

PiperOrigin-RevId: 797529699

Author: jnthntatum

tools/BUILD

@@ -81,7 +81,6 @@ cc_library(
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@com_google_cel_spec//proto/cel/expr:checked_cc_proto",
         "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",

tools/internal/BUILD

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -21,5 +22,18 @@ licenses(["notice"])
 cc_library(
     name = "navigable_ast_internal",
     hdrs = ["navigable_ast_internal.h"],
-    deps = ["@com_google_absl//absl/types:span"],
+    deps = [
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "navigable_ast_internal_test",
+    srcs = ["navigable_ast_internal_test.cc"],
+    deps = [
+        ":navigable_ast_internal",
+        "//internal:testing",
+        "@com_google_absl//absl/types:span",
+    ],
 )

tools/internal/navigable_ast_internal.h

@@ -11,10 +11,14 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-
 #ifndef THIRD_PARTY_CEL_CPP_TOOLS_INTERNAL_NAVIGABLE_AST_INTERNAL_H_
 #define THIRD_PARTY_CEL_CPP_TOOLS_INTERNAL_NAVIGABLE_AST_INTERNAL_H_
 
+#include <cstddef>
+#include <iterator>
+#include <type_traits>
+
+#include "absl/log/absl_check.h"
 #include "absl/types/span.h"
 
 namespace cel::tools_internal {
@@ -27,44 +31,76 @@ namespace cel::tools_internal {
 // RangeTraits provide type info needed to construct the span and adapt to the
 // range element type.
 template <class RangeTraits>
-class SpanRange {
+class NavigableAstRange {
  private:
   using UnderlyingType = typename RangeTraits::UnderlyingType;
+  using PtrType = const UnderlyingType*;
   using SpanType = absl::Span<const UnderlyingType>;
 
-  class SpanForwardIter {
+ public:
+  class Iterator {
    public:
-    SpanForwardIter(SpanType span, int i) : i_(i), span_(span) {}
+    using difference_type = ptrdiff_t;
+    using value_type = decltype(RangeTraits::Adapt(*PtrType()));
+    using iterator_category = std::bidirectional_iterator_tag;
+
+    Iterator() : ptr_(nullptr), span_() {}
+    Iterator(SpanType span, size_t i) : ptr_(span.data() + i), span_(span) {}
+
+    value_type operator*() const {
+      ABSL_DCHECK(ptr_ != nullptr);
+      ABSL_DCHECK(span_.data() != nullptr);
+      ABSL_DCHECK_GE(ptr_, span_.data());
+      ABSL_DCHECK_LT(ptr_, span_.data() + span_.size());
+      return RangeTraits::Adapt(*ptr_);
+    }
 
-    decltype(RangeTraits::Adapt(SpanType()[0])) operator*() const {
-      ABSL_CHECK(i_ < span_.size());
-      return RangeTraits::Adapt(span_[i_]);
+    template <int... Barrier, typename T = value_type>
+    std::enable_if_t<std::is_lvalue_reference<T>::value,
+                     std::add_pointer_t<std::remove_reference_t<T>>>
+    operator->() const {
+      return &operator*();
     }
 
-    SpanForwardIter& operator++() {
-      ++i_;
+    Iterator& operator++() {
+      ++ptr_;
       return *this;
     }
 
-    bool operator==(const SpanForwardIter& other) const {
-      return i_ == other.i_ && span_ == other.span_;
+    Iterator operator++(int) {
+      Iterator tmp = *this;
+      ++ptr_;
+      return tmp;
     }
 
-    bool operator!=(const SpanForwardIter& other) const {
-      return !(*this == other);
+    Iterator& operator--() {
+      --ptr_;
+      return *this;
     }
 
+    Iterator operator--(int) {
+      Iterator tmp = *this;
+      --ptr_;
+      return tmp;
+    }
+
+    bool operator==(const Iterator& other) const {
+      return ptr_ == other.ptr_ && span_ == other.span_;
+    }
+
+    bool operator!=(const Iterator& other) const { return !(*this == other); }
+
    private:
-    int i_;
+    PtrType ptr_;
     SpanType span_;
   };
 
- public:
-  explicit SpanRange(SpanType span) : span_(span) {}
+  explicit NavigableAstRange(SpanType span) : span_(span) {}
 
-  SpanForwardIter begin() { return SpanForwardIter(span_, 0); }
+  Iterator begin() { return Iterator(span_, 0); }
+  Iterator end() { return Iterator(span_, span_.size()); }
 
-  SpanForwardIter end() { return SpanForwardIter(span_, span_.size()); }
+  explicit operator bool() const { return !span_.empty(); }
 
  private:
   SpanType span_;

tools/internal/navigable_ast_internal_test.cc

@@ -0,0 +1,48 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "tools/internal/navigable_ast_internal.h"
+
+#include <iterator>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "internal/testing.h"
+
+namespace cel::tools_internal {
+namespace {
+
+struct TestRangeTraits {
+  using UnderlyingType = int;
+  static double Adapt(const UnderlyingType& value) {
+    return static_cast<double>(value) + 0.5;
+  }
+};
+
+TEST(NavigableAstRangeTest, BasicIteration) {
+  std::vector<int> values{1, 2, 3};
+  NavigableAstRange<TestRangeTraits> range(absl::MakeConstSpan(values));
+  absl::Span<const int> span(values);
+  auto it = range.begin();
+  EXPECT_EQ(*it, 1.5);
+  EXPECT_EQ(*++it, 2.5);
+  EXPECT_EQ(*++it, 3.5);
+  EXPECT_EQ(++it, range.end());
+  EXPECT_EQ(*--it, 3.5);
+  EXPECT_EQ(*--it, 2.5);
+  EXPECT_EQ(*--it, 1.5);
+  EXPECT_EQ(it, range.begin());
+}
+
+}  // namespace
+}  // namespace cel::tools_internal

tools/navigable_ast.cc

@@ -14,6 +14,7 @@
 
 #include "tools/navigable_ast.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <memory>
 #include <string>
@@ -143,21 +144,23 @@ class NavigableExprBuilderVisitor
                           ? nullptr
                           : metadata_->nodes[parent_stack_.back()].get();
     size_t index = metadata_->AddNode();
+    AstNode* node = metadata_->nodes[index].get();
     tools_internal::AstNodeData& node_data = metadata_->NodeDataAt(index);
     node_data.parent = parent;
     node_data.expr = expr;
     node_data.parent_relation = ChildKind::kUnspecified;
     node_data.node_kind = GetNodeKind(*expr);
-    node_data.weight = 1;
+    node_data.tree_size = 1;
+    node_data.height = 1;
     node_data.index = index;
     node_data.metadata = metadata_.get();
 
-    metadata_->id_to_node.insert({expr->id(), index});
-    metadata_->expr_to_node.insert({expr, index});
+    metadata_->id_to_node.insert({expr->id(), node});
+    metadata_->expr_to_node.insert({expr, node});
     if (!parent_stack_.empty()) {
       auto& parent_node_data = metadata_->NodeDataAt(parent_stack_.back());
       size_t child_index = parent_node_data.children.size();
-      parent_node_data.children.push_back(metadata_->nodes[index].get());
+      parent_node_data.children.push_back(node);
       node_data.parent_relation = GetChildKind(parent_node_data, child_index);
     }
     parent_stack_.push_back(index);
@@ -172,7 +175,9 @@ class NavigableExprBuilderVisitor
     if (!parent_stack_.empty()) {
       tools_internal::AstNodeData& parent_node_data =
           metadata_->NodeDataAt(parent_stack_.back());
-      parent_node_data.weight += node.weight;
+      parent_node_data.tree_size += node.tree_size;
+      parent_node_data.height =
+          std::max(parent_node_data.height, node.height + 1);
     }
   }
 
@@ -261,12 +266,12 @@ int AstNode::child_index() const {
 
 AstNode::PreorderRange AstNode::DescendantsPreorder() const {
   return AstNode::PreorderRange(absl::MakeConstSpan(data_.metadata->nodes)
-                                    .subspan(data_.index, data_.weight));
+                                    .subspan(data_.index, data_.tree_size));
 }
 
 AstNode::PostorderRange AstNode::DescendantsPostorder() const {
   return AstNode::PostorderRange(absl::MakeConstSpan(data_.metadata->postorder)
-                                     .subspan(data_.index, data_.weight));
+                                     .subspan(data_.index, data_.tree_size));
 }
 
 NavigableAst NavigableAst::Build(const Expr& expr) {

tools/navigable_ast.h

@@ -93,15 +93,21 @@ struct AstNodeData {
   NodeKind node_kind;
   const AstMetadata* metadata;
   size_t index;
-  size_t weight;
+  size_t tree_size;
+  size_t height;
   std::vector<AstNode*> children;
 };
 
 struct AstMetadata {
+  // The nodes in the AST in preorder.
+  //
+  // unique_ptr is used to guarantee pointer stability in the other tables.
   std::vector<std::unique_ptr<AstNode>> nodes;
-  std::vector<const AstNode*> postorder;
-  absl::flat_hash_map<int64_t, size_t> id_to_node;
-  absl::flat_hash_map<const cel::expr::Expr*, size_t> expr_to_node;
+  std::vector<const AstNode* absl_nonnull> postorder;
+  absl::flat_hash_map<int64_t, const AstNode* absl_nonnull> id_to_node;
+  absl::flat_hash_map<const cel::expr::Expr*,
+                      const AstNode* absl_nonnull>
+      expr_to_node;
 
   AstNodeData& NodeDataAt(size_t index);
   size_t AddNode();
@@ -123,13 +129,19 @@ struct PreorderTraits {
 
 // Wrapper around a CEL AST node that exposes traversal information.
 class AstNode {
- private:
+ public:
+  // A const Span like type that provides pre-order traversal for a sub tree.
+  // provides .begin() and .end() returning bidirectional iterators to
+  // const AstNode&.
   using PreorderRange =
-      tools_internal::SpanRange<tools_internal::PreorderTraits>;
+      tools_internal::NavigableAstRange<tools_internal::PreorderTraits>;
+
+  // A const Span like type that provides post-order traversal for a sub tree.
+  // provides .begin() and .end() returning bidirectional iterators to
+  // const AstNode&.
   using PostorderRange =
-      tools_internal::SpanRange<tools_internal::PostorderTraits>;
+      tools_internal::NavigableAstRange<tools_internal::PostorderTraits>;
 
- public:
   // The parent of this node or nullptr if it is a root.
   const AstNode* absl_nullable parent() const { return data_.parent; }
 
@@ -146,6 +158,13 @@ class AstNode {
   // The type of this node, analogous to Expr::ExprKindCase.
   NodeKind node_kind() const { return data_.node_kind; }
 
+  // The number of nodes in the tree rooted at this node (including self).
+  size_t tree_size() const { return data_.tree_size; }
+
+  // The height of this node in the tree (the number of descendants including
+  // self on the longest path).
+  size_t height() const { return data_.height; }
+
   absl::Span<const AstNode* const> children() const {
     return absl::MakeConstSpan(data_.children);
   }
@@ -164,9 +183,6 @@ class AstNode {
   //   - maps are traversed in order (alternating key, value per entry)
   //   - comprehensions are traversed in the order: range, accu_init, condition,
   //   step, result
-  //
-  // Return type is an implementation detail, it should only be used with auto
-  // or in a range-for loop.
   PreorderRange DescendantsPreorder() const;
 
   // Range over the descendants of this node (including self) using postorder
@@ -184,12 +200,16 @@ class AstNode {
 };
 
 // NavigableExpr provides a view over a CEL AST that allows for generalized
-// traversal.
+// traversal. The traversal structures are eagerly built on construction,
+// requiring a full traversal of the AST. This is intended for use in tools that
+// might require random access or multiple passes over the AST, amortizing the
+// cost of building the traversal structures.
 //
 // Pointers to AstNodes are owned by this instance and must not outlive it.
 //
-// Note: Assumes ptr stability of the input Expr pb -- this is only guaranteed
-// if no mutations take place on the input.
+// `NavigableAst` and Navigable nodes are independent of the input Expr and may
+// outlive it, but may contain dangling pointers if the input Expr is modified
+// or destroyed.
 class NavigableAst {
  public:
   static NavigableAst Build(const cel::expr::Expr& expr);
@@ -217,7 +237,7 @@ class NavigableAst {
     if (it == metadata_->id_to_node.end()) {
       return nullptr;
     }
-    return metadata_->nodes[it->second].get();
+    return it->second;
   }
 
   // Return ptr to the AST node representing the given Expr protobuf node.
@@ -227,7 +247,7 @@ class NavigableAst {
     if (it == metadata_->expr_to_node.end()) {
       return nullptr;
     }
-    return metadata_->nodes[it->second].get();
+    return it->second;
   }
 
   // The root of the AST.