stackless traversal

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" }

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
     }
 }

src/algo/mod.rs

@@ -0,0 +1 @@
+pub mod walk;

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)
+    })
+}
+

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)
-        }
-    }
 }

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},

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);

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)
+    }
 }

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() }
     }