Add TreeRecursiveWalker

Adds a method of breadth-first recursive descent for TreeCursors
pull/6198/head
Skyler Hawthorne 8 months ago
parent 24eb54b2d0
commit cc5e1937c4

@ -1,7 +1,8 @@
use std::{cmp::Reverse, ops::Range}; use std::{cmp::Reverse, collections::VecDeque, ops::Range};
use super::{LanguageLayer, LayerId}; use super::{LanguageLayer, LayerId};
use ropey::RopeSlice;
use slotmap::HopSlotMap; use slotmap::HopSlotMap;
use tree_sitter::Node; use tree_sitter::Node;
@ -18,18 +19,18 @@ struct InjectionRange {
depth: u32, depth: u32,
} }
pub struct TreeCursor<'a> { pub struct TreeCursor<'n> {
layers: &'a HopSlotMap<LayerId, LanguageLayer>, layers: &'n HopSlotMap<LayerId, LanguageLayer>,
root: LayerId, root: LayerId,
current: LayerId, current: LayerId,
injection_ranges: Vec<InjectionRange>, injection_ranges: Vec<InjectionRange>,
// TODO: Ideally this would be a `tree_sitter::TreeCursor<'a>` but // TODO: Ideally this would be a `tree_sitter::TreeCursor<'a>` but
// that returns very surprising results in testing. // that returns very surprising results in testing.
cursor: Node<'a>, cursor: Node<'n>,
} }
impl<'a> TreeCursor<'a> { impl<'n> TreeCursor<'n> {
pub(super) fn new(layers: &'a HopSlotMap<LayerId, LanguageLayer>, root: LayerId) -> Self { pub(super) fn new(layers: &'n HopSlotMap<LayerId, LanguageLayer>, root: LayerId) -> Self {
let mut injection_ranges = Vec::new(); let mut injection_ranges = Vec::new();
for (layer_id, layer) in layers.iter() { 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 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
}
}

Loading…
Cancel
Save