From b2d46b508b733fa8e26c524c3ebc022e6a514064 Mon Sep 17 00:00:00 2001
From: asoffer <andy.soffer@gmail.com>
Date: Fri, 17 Nov 2023 18:01:21 -0500
Subject: [PATCH] Ancestor iteration for LexicalScopes

---
 ir/BUILD            |  1 +
 ir/lexical_scope.cc | 11 ++------
 ir/lexical_scope.h  | 67 +++++++++++++++++++++++++++++++++++++++++----
 3 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/ir/BUILD b/ir/BUILD
index c8600484..7edc487a 100644
--- a/ir/BUILD
+++ b/ir/BUILD
@@ -208,6 +208,7 @@ cc_library(
         "//parse:node_index",
         "//type",
         "@asoffer_nth//nth/debug",
+        "@asoffer_nth//nth/utility:iterator_range",
         "@com_google_absl//absl/container:flat_hash_map",
     ]
 )
diff --git a/ir/lexical_scope.cc b/ir/lexical_scope.cc
index 116cbaa2..0d95fb07 100644
--- a/ir/lexical_scope.cc
+++ b/ir/lexical_scope.cc
@@ -10,8 +10,6 @@ LexicalScope::Index LexicalScope::Index::Invalid() {
   return Index(std::numeric_limits<Index::underlying_type>::max());
 }
 
-LexicalScope::Index LexicalScope::parent() const { return parent_; }
-
 LexicalScope::DeclarationInfo const *LexicalScope::identifier(
     Identifier id) const {
   auto iter = identifiers_.find(id);
@@ -22,7 +20,7 @@ LexicalScope::DeclarationInfo const *LexicalScope::identifier(
 LexicalScope::Index LexicalScopeTree::insert_child(
     LexicalScope::Index parent_index) {
   LexicalScope::Index index(scopes_.size());
-  scopes_.push_back(LexicalScope(parent_index));
+  scopes_.push_back(LexicalScope(index.value() - parent_index.value()));
   return index;
 }
 
@@ -42,11 +40,8 @@ LexicalScope const &LexicalScopeTree::operator[](
 
 LexicalScope::DeclarationInfo const *LexicalScopeTree::identifier(
     LexicalScope::Index index, Identifier id) const {
-  while (index != LexicalScope::Index::Invalid()) {
-    auto const &scope = scopes_[index.value()];
-    auto const *info  = scope.identifier(id);
-    if (info) { return info; }
-    index = scope.parent();
+  for (auto const &scope : ancestors(index)) {
+    if (auto const *info = scope.identifier(id)) { return info; }
   }
   return nullptr;
 }
diff --git a/ir/lexical_scope.h b/ir/lexical_scope.h
index dbe222ec..e19b396a 100644
--- a/ir/lexical_scope.h
+++ b/ir/lexical_scope.h
@@ -7,6 +7,7 @@
 #include "common/identifier.h"
 #include "common/strong_identifier_type.h"
 #include "nth/debug/debug.h"
+#include "nth/utility/iterator_range.h"
 #include "parse/node_index.h"
 #include "type/type.h"
 
@@ -22,8 +23,6 @@ struct LexicalScope {
     static Index Invalid();
   };
 
-  Index parent() const;
-
   struct DeclarationInfo {
     ParseNodeIndex declaration;
     ParseNodeIndex identifier;
@@ -40,11 +39,12 @@ struct LexicalScope {
  private:
   friend struct LexicalScopeTree;
 
-  explicit LexicalScope(Index parent_index) : parent_(parent_index) {}
+  explicit LexicalScope(uint32_t distance_to_parent)
+      : parent_distance_(distance_to_parent) {}
 
   absl::flat_hash_map<Identifier, DeclarationInfo> identifiers_;
 
-  Index parent_;
+  uint32_t parent_distance_;
 };
 
 struct LexicalScopeTree {
@@ -53,16 +53,71 @@ struct LexicalScopeTree {
   LexicalScope &root();
   LexicalScope const &root() const;
 
+ private:
+  struct ancestor_iterator_base {
+    friend bool operator==(ancestor_iterator_base,
+                           ancestor_iterator_base) = default;
+
+   protected:
+    explicit ancestor_iterator_base(LexicalScope const *scope)
+        : scope_(scope) {}
+
+    LexicalScope const *scope_;
+  };
+
+ public:
+  template <typename S>
+  struct ancestor_iterator_impl : ancestor_iterator_base {
+    ancestor_iterator_impl &operator++() {
+      LexicalScopeTree::ReplaceWithParent(scope_);
+      return *this;
+    }
+
+    ancestor_iterator_impl operator++(int) {
+      auto copy = *this;
+      ++*this;
+      return copy;
+    }
+
+    S &operator*() const { return *const_cast<S *>(scope_); }
+    S *operator->() const { return const_cast<S *>(scope_); }
+
+   private:
+    friend struct LexicalScopeTree;
+    using ancestor_iterator_base::ancestor_iterator_base;
+  };
+
+  using ancestor_iterator = ancestor_iterator_impl<LexicalScope>;
+  using const_ancestor_iterator = ancestor_iterator_impl<LexicalScope const>;
+
+  auto ancestors(LexicalScope::Index index) const {
+    return nth::iterator_range(
+        const_ancestor_iterator(scopes_.data() + index.value()),
+        const_ancestor_iterator(scopes_.data() - 1));
+  }
+
+  auto ancestors(LexicalScope::Index index) {
+    return nth::iterator_range(
+        ancestor_iterator(scopes_.data() + index.value()),
+        ancestor_iterator(scopes_.data()));
+  }
+
   LexicalScope &operator[](LexicalScope::Index index);
   LexicalScope const &operator[](LexicalScope::Index index) const;
 
   LexicalScope::DeclarationInfo const *identifier(LexicalScope::Index index,
                                                   Identifier id) const;
 
-  LexicalScopeTree()
-      : scopes_(1, LexicalScope(LexicalScope::Index::Invalid())) {}
+  LexicalScopeTree() : scopes_(1, LexicalScope(1)) {}
 
  private:
+  template <typename>
+  friend struct ancestor_iterator_impl;
+
+  static void ReplaceWithParent(LexicalScope const *&scope) {
+    scope -= scope->parent_distance_;
+  }
+
   std::vector<LexicalScope> scopes_;
 };