1 files changed, 51 insertions, 0 deletions

diff --git a/src/traverse.rs b/src/traverse.rs
new file mode 100644
index 0000000..df13f5d
--- /dev/null
+++ b/src/traverse.rs
@@ -0,0 +1,51 @@
+// Lazy pre-order traversal over arbitrary tree-sitter trees
+
+// extern
+use tree_sitter::{Node, Tree, TreeCursor};
+
+pub struct PreOrder<'a> {
+    cursor: TreeCursor<'a>,
+    done: bool,
+}
+
+impl<'a> PreOrder<'a> {
+    pub fn new(tree: &'a Tree) -> Self {
+        Self {
+            cursor: tree.walk(),
+            done: false,
+        }
+    }
+}
+
+// Traverse the tree in root-left-right fashion
+impl<'a> Iterator for PreOrder<'a> {
+    type Item = Node<'a>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done {
+            return None;
+        }
+
+        // capture the root node
+        let node = self.cursor.node();
+
+        // if this node has a child, move there and return the root, or
+        // if this node has a sibling, move there and return the root
+        if self.cursor.goto_first_child() || self.cursor.goto_next_sibling() {
+            return Some(node);
+        }
+
+        loop {
+            // we are done retracing all the way back to the root node
+            if !self.cursor.goto_parent() {
+                self.done = true;
+                break;
+            }
+
+            if self.cursor.goto_next_sibling() {
+                break;
+            }
+        }
+        Some(node)
+    }
+}