From cc5e1937c4b45dceddc7b91912b478642e0dc604 Mon Sep 17 00:00:00 2001
From: Skyler Hawthorne <skyler@dead10ck.com>
Date: Sun, 7 Apr 2024 14:26:52 -0400
Subject: [PATCH] Add TreeRecursiveWalker

Adds a method of breadth-first recursive descent for TreeCursors
---
 helix-core/src/syntax/tree_cursor.rs | 66 +++++++++++++++++++++++++---
 1 file changed, 59 insertions(+), 7 deletions(-)

diff --git a/helix-core/src/syntax/tree_cursor.rs b/helix-core/src/syntax/tree_cursor.rs
index bec4a1c6c..cfcb3839f 100644
--- a/helix-core/src/syntax/tree_cursor.rs
+++ b/helix-core/src/syntax/tree_cursor.rs
@@ -1,7 +1,8 @@
-use std::{cmp::Reverse, ops::Range};
+use std::{cmp::Reverse, collections::VecDeque, ops::Range};
 
 use super::{LanguageLayer, LayerId};
 
+use ropey::RopeSlice;
 use slotmap::HopSlotMap;
 use tree_sitter::Node;
 
@@ -18,18 +19,18 @@ struct InjectionRange {
     depth: u32,
 }
 
-pub struct TreeCursor<'a> {
-    layers: &'a HopSlotMap<LayerId, LanguageLayer>,
+pub struct TreeCursor<'n> {
+    layers: &'n HopSlotMap<LayerId, LanguageLayer>,
     root: LayerId,
     current: LayerId,
     injection_ranges: Vec<InjectionRange>,
     // TODO: Ideally this would be a `tree_sitter::TreeCursor<'a>` but
     // that returns very surprising results in testing.
-    cursor: Node<'a>,
+    cursor: Node<'n>,
 }
 
-impl<'a> TreeCursor<'a> {
-    pub(super) fn new(layers: &'a HopSlotMap<LayerId, LanguageLayer>, root: LayerId) -> Self {
+impl<'n> TreeCursor<'n> {
+    pub(super) fn new(layers: &'n HopSlotMap<LayerId, LanguageLayer>, root: LayerId) -> Self {
         let mut injection_ranges = Vec::new();
 
         for (layer_id, layer) in layers.iter() {
@@ -61,7 +62,7 @@ impl<'a> TreeCursor<'a> {
         }
     }
 
-    pub fn node(&self) -> Node<'a> {
+    pub fn node(&self) -> Node<'n> {
         self.cursor
     }
 
@@ -262,3 +263,54 @@ impl<'n> Iterator for ChildIter<'n> {
         }
     }
 }
+
+impl<'n> IntoIterator for &'n mut TreeCursor<'n> {
+    type Item = Node<'n>;
+    type IntoIter = TreeRecursiveWalker<'n>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        let mut queue = VecDeque::new();
+        let root = self.node();
+        queue.push_back(root);
+
+        TreeRecursiveWalker {
+            cursor: self,
+            queue,
+            root,
+        }
+    }
+}
+
+pub struct TreeRecursiveWalker<'n> {
+    cursor: &'n mut TreeCursor<'n>,
+    queue: VecDeque<Node<'n>>,
+    root: Node<'n>,
+}
+
+impl<'n> Iterator for TreeRecursiveWalker<'n> {
+    type Item = Node<'n>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let current = self.cursor.node();
+        log::debug!("recursive walk -- current: {current:?}");
+
+        if current != self.root && self.cursor.goto_next_named_sibling() {
+            self.queue.push_back(current);
+            log::debug!("recursive walk -- sibling: {:?}", self.cursor.node());
+            return Some(self.cursor.node());
+        }
+
+        while let Some(queued) = self.queue.pop_front() {
+            self.cursor.cursor = queued;
+
+            if !self.cursor.goto_first_named_child() {
+                continue;
+            }
+
+            log::debug!("recursive walk -- child: {:?}", self.cursor.node());
+            return Some(self.cursor.node());
+        }
+
+        None
+    }
+}