stackless traversal

9 files changed, 86 insertions, 16 deletions

diff --git a/Cargo.toml b/Cargo.toml
index b89573e2b..e0c1ce0a4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -10,6 +10,7 @@ members = [ "tools", "cli" ]
 [dependencies]
 unicode-xid = "0.1.0"
 text_unit = "0.1.1"
+itertools = "0.7.5"
 
 [dev-dependencies]
 testutils = { path = "./tests/testutils" }
diff --git a/code/native/src/lib.rs b/code/native/src/lib.rs
index dcf478cf5..068767fab 100644
--- a/code/native/src/lib.rs
+++ b/code/native/src/lib.rs
@@ -7,6 +7,7 @@ use libsyntax2::{
     File,
     utils::dump_tree,
     SyntaxKind::*,
+    algo,
 };
 use neon::prelude::*;
 
@@ -17,11 +18,12 @@ pub struct Wrapper {
 impl Wrapper {
     fn highlight(&self) -> Vec<(TextRange, &'static str)> {
         let mut res = Vec::new();
-        self.inner.for_each_node(|node| {
+        let syntax = self.inner.syntax();
+        for node in algo::walk::preorder(syntax.as_ref()) {
             if node.kind() == ERROR {
                 res.push((node.range(), "error"))
             }
-        });
+        }
         res
     }
 }
diff --git a/src/algo/mod.rs b/src/algo/mod.rs
new file mode 100644
index 000000000..c1644d16d
--- /dev/null
+++ b/src/algo/mod.rs
@@ -0,0 +1 @@
+pub mod walk;
diff --git a/src/algo/walk.rs b/src/algo/walk.rs
new file mode 100644
index 000000000..86dd82cc9
--- /dev/null
+++ b/src/algo/walk.rs
@@ -0,0 +1,46 @@
+use SyntaxNodeRef;
+
+pub fn preorder<'a>(root: SyntaxNodeRef<'a>) -> impl Iterator<Item=SyntaxNodeRef<'a>> {
+    walk(root).filter_map(|event| match event {
+        WalkEvent::Enter(node) => Some(node),
+        WalkEvent::Exit(_) => None,
+    })
+}
+
+#[derive(Debug, Copy, Clone)]
+enum WalkEvent<'a> {
+    Enter(SyntaxNodeRef<'a>),
+    Exit(SyntaxNodeRef<'a>),
+}
+
+fn walk<'a>(root: SyntaxNodeRef<'a>) -> impl Iterator<Item=WalkEvent<'a>> {
+    let mut done = false;
+    ::itertools::unfold(WalkEvent::Enter(root), move |pos| {
+        if done {
+            return None;
+        }
+        let res = *pos;
+        *pos = match *pos {
+            WalkEvent::Enter(node) => match node.first_child() {
+                Some(child) => WalkEvent::Enter(child),
+                None => WalkEvent::Exit(node),
+            },
+            WalkEvent::Exit(node) => {
+                if node == root {
+                    done = true;
+                    WalkEvent::Exit(node)
+                } else {
+                    match node.next_sibling() {
+                        Some(sibling) => WalkEvent::Enter(sibling),
+                        None => match node.parent() {
+                            Some(node) => WalkEvent::Exit(node),
+                            None => WalkEvent::Exit(node),
+                        }
+                    }
+                }
+            }
+        };
+        Some(res)
+    })
+}
+
diff --git a/src/ast.rs b/src/ast.rs
index 3a9287466..b513eb13e 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,5 +1,5 @@
 use std::sync::Arc;
-use {SyntaxNode, TreeRoot, SyntaxRoot, SyntaxNodeRef};
+use {SyntaxNode, TreeRoot, SyntaxRoot};
 
 #[derive(Debug)]
 pub struct File<R: TreeRoot = Arc<SyntaxRoot>> {
@@ -16,15 +16,4 @@ impl<R: TreeRoot> File<R> {
     pub fn syntax(&self) -> SyntaxNode<R> {
         self.syntax.clone()
     }
-
-    pub fn for_each_node(&self, mut f: impl FnMut(SyntaxNodeRef)) {
-        let syntax = self.syntax();
-        let syntax = syntax.borrow();
-        go(syntax, &mut f);
-
-        fn go(syntax: SyntaxNodeRef, f: &mut FnMut(SyntaxNodeRef)) {
-            f(syntax);
-            syntax.children().into_iter().for_each(f)
-        }
-    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 3a8e6fa64..9049beb29 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -22,6 +22,7 @@
 
 extern crate text_unit;
 extern crate unicode_xid;
+extern crate itertools;
 
 mod lexer;
 mod parser;
@@ -30,6 +31,7 @@ mod yellow;
 /// Utilities for simple uses of the parser.
 pub mod utils;
 pub mod ast;
+pub mod algo;
 
 pub use {
     lexer::{tokenize, Token},
diff --git a/src/utils.rs b/src/utils.rs
index 20991659a..826a7d60b 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -3,7 +3,7 @@ use {SyntaxError, SyntaxNode, SyntaxNodeRef};
 
 /// Parse a file and create a string representation of the resulting parse tree.
 pub fn dump_tree(syntax: &SyntaxNode) -> String {
-    let syntax = syntax.borrow();
+    let syntax = syntax.as_ref();
     let mut errors: BTreeSet<_> = syntax.root.errors.iter().cloned().collect();
     let mut result = String::new();
     go(syntax, &mut result, 0, &mut errors);
diff --git a/src/yellow/red.rs b/src/yellow/red.rs
index bffb72510..8329ec5b2 100644
--- a/src/yellow/red.rs
+++ b/src/yellow/red.rs
@@ -84,4 +84,7 @@ impl RedNode {
     pub(crate) fn parent(&self) -> Option<ptr::NonNull<RedNode>> {
         Some(self.parent.as_ref()?.parent)
     }
+    pub(crate) fn index_in_parent(&self) -> Option<usize> {
+        Some(self.parent.as_ref()?.index_in_parent)
+    }
 }
diff --git a/src/yellow/syntax.rs b/src/yellow/syntax.rs
index 5c31a3f35..65ce647c7 100644
--- a/src/yellow/syntax.rs
+++ b/src/yellow/syntax.rs
@@ -18,6 +18,15 @@ pub struct SyntaxNode<R: TreeRoot = Arc<SyntaxRoot>> {
     red: ptr::NonNull<RedNode>,
 }
 
+impl <R1: TreeRoot, R2: TreeRoot> PartialEq<SyntaxNode<R1>> for SyntaxNode<R2> {
+    fn eq(&self, other: &SyntaxNode<R1>) -> bool {
+        self.red == other.red
+    }
+}
+
+impl <R: TreeRoot> Eq for SyntaxNode<R> {
+}
+
 pub type SyntaxNodeRef<'a> = SyntaxNode<&'a SyntaxRoot>;
 
 #[derive(Debug)]
@@ -53,7 +62,7 @@ impl SyntaxNode<Arc<SyntaxRoot>> {
 }
 
 impl<R: TreeRoot> SyntaxNode<R> {
-    pub fn borrow<'a>(&'a self) -> SyntaxNode<&'a SyntaxRoot> {
+    pub fn as_ref<'a>(&'a self) -> SyntaxNode<&'a SyntaxRoot> {
         SyntaxNode {
             root: &*self.root,
             red: ptr::NonNull::clone(&self.red),
@@ -92,6 +101,23 @@ impl<R: TreeRoot> SyntaxNode<R> {
         })
     }
 
+    pub fn first_child(&self) -> Option<SyntaxNode<R>> {
+        self.children().next()
+    }
+
+    pub fn next_sibling(&self) -> Option<SyntaxNode<R>> {
+        let red = self.red();
+        let parent = self.parent()?;
+        let next_sibling_idx = red.index_in_parent()? + 1;
+        if next_sibling_idx == red.n_children() {
+            return None;
+        }
+        Some(SyntaxNode {
+            root: self.root.clone(),
+            red: parent.red().nth_child(next_sibling_idx),
+        })
+    }
+
     fn red(&self) -> &RedNode {
         unsafe { self.red.as_ref() }
     }