internal change

tools/BUILD

@@ -46,10 +46,12 @@ cc_library(
         "//eval/public:ast_visitor",
         "//eval/public:ast_visitor_base",
         "//eval/public:source_position",
+        "//tools/internal:navigable_ast_internal",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@com_google_googleapis//google/api/expr/v1alpha1:checked_cc_proto",
         "@com_google_googleapis//google/api/expr/v1alpha1:syntax_cc_proto",

tools/internal/BUILD

@@ -0,0 +1,23 @@
+# Copyright 2023 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.
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])
+
+cc_library(
+    name = "navigable_ast_internal",
+    hdrs = ["navigable_ast_internal.h"],
+    deps = ["@com_google_absl//absl/types:span"],
+)

tools/internal/navigable_ast_internal.h

@@ -0,0 +1,75 @@
+// Copyright 2023 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.
+
+#ifndef THIRD_PARTY_CEL_CPP_TOOLS_INTERNAL_NAVIGABLE_AST_INTERNAL_H_
+#define THIRD_PARTY_CEL_CPP_TOOLS_INTERNAL_NAVIGABLE_AST_INTERNAL_H_
+
+#include "absl/types/span.h"
+
+namespace cel::tools_internal {
+
+// Implementation for range used for traversals backed by an absl::Span.
+//
+// This is intended to abstract the metadata layout from clients using the
+// traversal methods in navigable_expr.h
+//
+// RangeTraits provide type info needed to construct the span and adapt to the
+// range element type.
+template <class RangeTraits>
+class SpanRange {
+ private:
+  using UnderlyingType = typename RangeTraits::UnderlyingType;
+  using SpanType = absl::Span<const UnderlyingType>;
+
+  class SpanForwardIter {
+   public:
+    SpanForwardIter(SpanType span, int i) : i_(i), span_(span) {}
+
+    auto operator*() const {
+      ABSL_CHECK(i_ < span_.size());
+      return RangeTraits::Adapt(span_[i_]);
+    }
+
+    SpanForwardIter& operator++() {
+      ++i_;
+      return *this;
+    }
+
+    bool operator==(const SpanForwardIter& other) const {
+      return i_ == other.i_ && span_ == other.span_;
+    }
+
+    bool operator!=(const SpanForwardIter& other) const {
+      return !(*this == other);
+    }
+
+   private:
+    int i_;
+    SpanType span_;
+  };
+
+ public:
+  explicit SpanRange(SpanType span) : span_(span) {}
+
+  SpanForwardIter begin() { return SpanForwardIter(span_, 0); }
+
+  SpanForwardIter end() { return SpanForwardIter(span_, span_.size()); }
+
+ private:
+  SpanType span_;
+};
+
+}  // namespace cel::tools_internal
+
+#endif  // THIRD_PARTY_CEL_CPP_TOOLS_INTERNAL_NAVIGABLE_AST_INTERNAL_H_

tools/navigable_ast.cc

@@ -23,14 +23,16 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/absl_check.h"
 #include "absl/memory/memory.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "eval/public/ast_traverse.h"
 #include "eval/public/ast_visitor.h"
 #include "eval/public/ast_visitor_base.h"
 #include "eval/public/source_position.h"
 
 namespace cel {
 
-namespace internal {
+namespace tools_internal {
 
 AstNodeData& AstMetadata::NodeDataAt(size_t index) {
   ABSL_CHECK(index < nodes.size());
@@ -43,7 +45,7 @@ size_t AstMetadata::AddNode() {
   return index;
 }
 
-}  // namespace internal
+}  // namespace tools_internal
 
 namespace {
 
@@ -79,7 +81,7 @@ NodeKind GetNodeKind(const Expr& expr) {
 
 // Get the traversal relationship from parent to the given node.
 // Note: these depend on the ast_visitor utility's traversal ordering.
-ChildKind GetChildKind(const internal::AstNodeData& parent_node,
+ChildKind GetChildKind(const tools_internal::AstNodeData& parent_node,
                        size_t child_index) {
   constexpr size_t kComprehensionRangeArgIndex =
       google::api::expr::runtime::ITER_RANGE;
@@ -112,9 +114,9 @@ ChildKind GetChildKind(const internal::AstNodeData& parent_node,
     case NodeKind::kComprehension:
       switch (child_index) {
         case kComprehensionRangeArgIndex:
-          return ChildKind::kCompehensionRange;
+          return ChildKind::kComprehensionRange;
         case kComprehensionInitArgIndex:
-          return ChildKind::kCompehensionInit;
+          return ChildKind::kComprehensionInit;
         case kComprehensionConditionArgIndex:
           return ChildKind::kComprehensionCondition;
         case kComprehensionLoopStepArgIndex:
@@ -133,18 +135,21 @@ class NavigableExprBuilderVisitor
     : public google::api::expr::runtime::AstVisitorBase {
  public:
   NavigableExprBuilderVisitor()
-      : metadata_(std::make_unique<internal::AstMetadata>()) {}
+      : metadata_(std::make_unique<tools_internal::AstMetadata>()) {}
 
   void PreVisitExpr(const Expr* expr, const SourcePosition* position) override {
     AstNode* parent = parent_stack_.empty()
                           ? nullptr
                           : metadata_->nodes[parent_stack_.back()].get();
     size_t index = metadata_->AddNode();
-    internal::AstNodeData& node_data = metadata_->NodeDataAt(index);
+    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.index = index;
+    node_data.metadata = metadata_.get();
 
     metadata_->id_to_node.insert({expr->id(), index});
     metadata_->expr_to_node.insert({expr, index});
@@ -159,20 +164,86 @@ class NavigableExprBuilderVisitor
 
   void PostVisitExpr(const Expr* expr,
                      const SourcePosition* position) override {
+    size_t idx = parent_stack_.back();
     parent_stack_.pop_back();
+    metadata_->postorder.push_back(metadata_->nodes[idx].get());
+    tools_internal::AstNodeData& node = metadata_->NodeDataAt(idx);
+    if (!parent_stack_.empty()) {
+      tools_internal::AstNodeData& parent_node_data =
+          metadata_->NodeDataAt(parent_stack_.back());
+      parent_node_data.weight += node.weight;
+    }
   }
 
-  std::unique_ptr<internal::AstMetadata> Consume() && {
+  std::unique_ptr<tools_internal::AstMetadata> Consume() && {
     return std::move(metadata_);
   }
 
  private:
-  std::unique_ptr<internal::AstMetadata> metadata_;
+  std::unique_ptr<tools_internal::AstMetadata> metadata_;
   std::vector<size_t> parent_stack_;
 };
 
 }  // namespace
 
+absl::string_view ChildKindName(ChildKind kind) {
+  switch (kind) {
+    case ChildKind::kUnspecified:
+      return "Unspecified";
+    case ChildKind::kSelectOperand:
+      return "SelectOperand";
+    case ChildKind::kCallReceiver:
+      return "CallReceiver";
+    case ChildKind::kCallArg:
+      return "CallArg";
+    case ChildKind::kListElem:
+      return "ListElem";
+    case ChildKind::kMapKey:
+      return "MapKey";
+    case ChildKind::kMapValue:
+      return "MapValue";
+    case ChildKind::kStructValue:
+      return "StructValue";
+    case ChildKind::kComprehensionRange:
+      return "CompehensionRange";
+    case ChildKind::kComprehensionInit:
+      return "CompehensionInit";
+    case ChildKind::kComprehensionCondition:
+      return "CompehensionCondition";
+    case ChildKind::kComprehensionLoopStep:
+      return "CompehensionLoopStep";
+    case ChildKind::kComprensionResult:
+      return "ComprensionResult";
+    default:
+      return "Unknown ChildKind";
+  }
+}
+
+absl::string_view NodeKindName(NodeKind kind) {
+  switch (kind) {
+    case NodeKind::kUnspecified:
+      return "Unspecified";
+    case NodeKind::kConstant:
+      return "Constant";
+    case NodeKind::kIdent:
+      return "Ident";
+    case NodeKind::kSelect:
+      return "Select";
+    case NodeKind::kCall:
+      return "Call";
+    case NodeKind::kList:
+      return "List";
+    case NodeKind::kMap:
+      return "Map";
+    case NodeKind::kStruct:
+      return "Struct";
+    case NodeKind::kComprehension:
+      return "Comprehension";
+    default:
+      return "Unknown NodeKind";
+  }
+}
+
 int AstNode::child_index() const {
   if (data_.parent == nullptr) {
     return -1;
@@ -187,6 +258,16 @@ int AstNode::child_index() const {
   return -1;
 }
 
+AstNode::PreorderRange AstNode::DescendantsPreorder() const {
+  return AstNode::PreorderRange(absl::MakeConstSpan(data_.metadata->nodes)
+                                    .subspan(data_.index, data_.weight));
+}
+
+AstNode::PostorderRange AstNode::DescendantsPostorder() const {
+  return AstNode::PostorderRange(absl::MakeConstSpan(data_.metadata->postorder)
+                                     .subspan(data_.index, data_.weight));
+}
+
 NavigableAst NavigableAst::Build(const Expr& expr) {
   NavigableExprBuilderVisitor visitor;
   AstTraverse(&expr, /*source_info=*/nullptr, &visitor);