5 files changed, 143 insertions, 225 deletions

diff --git a/src/yellow/builder.rs b/src/yellow/builder.rs
index 346d561cd..73b4b5ff9 100644
--- a/src/yellow/builder.rs
+++ b/src/yellow/builder.rs
@@ -1,13 +1,12 @@
-use std::sync::Arc;
 use {
     SyntaxKind, TextRange, TextUnit,
-    yellow::{SyntaxNode, GreenNode, SyntaxError},
+    yellow::{SyntaxNode, GreenNode, GreenNodeBuilder, SyntaxError},
     parser::Sink
 };
 
 pub(crate) struct GreenBuilder {
     text: String,
-    stack: Vec<GreenNode>,
+    stack: Vec<GreenNodeBuilder>,
     pos: TextUnit,
     root: Option<GreenNode>,
     errors: Vec<SyntaxError>,
@@ -33,24 +32,21 @@ impl Sink for GreenBuilder {
     fn leaf(&mut self, kind: SyntaxKind, len: TextUnit) {
         let range = TextRange::offset_len(self.pos, len);
         self.pos += len;
-        let text = self.text[range].to_owned();
+        let text = &self.text[range];
+        let leaf = GreenNodeBuilder::new_leaf(kind, text);
         let parent = self.stack.last_mut().unwrap();
-        if kind.is_trivia() {
-            parent.push_trivia(kind, text);
-        } else {
-            let node = GreenNode::new_leaf(kind, text);
-            parent.push_child(Arc::new(node));
-        }
+        parent.push_child(leaf)
     }
 
     fn start_internal(&mut self, kind: SyntaxKind) {
-        self.stack.push(GreenNode::new_branch(kind))
+        self.stack.push(GreenNodeBuilder::new_internal(kind))
     }
 
     fn finish_internal(&mut self) {
-        let node = self.stack.pop().unwrap();
+        let builder = self.stack.pop().unwrap();
+        let node = builder.build();
         if let Some(parent) = self.stack.last_mut() {
-            parent.push_child(Arc::new(node))
+            parent.push_child(node);
         } else {
             self.root = Some(node);
         }
@@ -61,7 +57,7 @@ impl Sink for GreenBuilder {
     }
 
     fn finish(self) -> SyntaxNode {
-        SyntaxNode::new(Arc::new(self.root.unwrap()), self.errors)
+        SyntaxNode::new(self.root.unwrap(), self.errors)
     }
 }
 
diff --git a/src/yellow/green.rs b/src/yellow/green.rs
index dafe1bb22..cb9dff128 100644
--- a/src/yellow/green.rs
+++ b/src/yellow/green.rs
@@ -1,187 +1,159 @@
 use std::sync::Arc;
-use text_unit::TextUnit;
-use SyntaxKind;
+use {SyntaxKind::{self, *}, TextUnit};
 
-type TokenText = String;
-
-#[derive(Debug)]
-pub(crate) struct GreenNode {
-    kind: SyntaxKind,
-    data: GreenNodeData,
+#[derive(Clone, Debug)]
+pub(crate) enum GreenNode {
+    Leaf(GreenLeaf),
+    Branch(Arc<GreenBranch>),
 }
 
 impl GreenNode {
-    pub(crate) fn new_leaf(kind: SyntaxKind, text: TokenText) -> GreenNode {
-        GreenNode {
-            kind,
-            data: GreenNodeData::Leaf(GreenLeaf { text }),
+    pub fn kind(&self) -> SyntaxKind {
+        match self {
+            GreenNode::Leaf(l) => l.kind(),
+            GreenNode::Branch(b) => b.kind(),
         }
     }
 
-    pub(crate) fn new_branch(
-        kind: SyntaxKind,
-    ) -> GreenNode {
-        let branch = GreenBranch {
-            text_len: 0.into(),
-            leading_trivia: Trivias::default(),
-            children: Vec::new(),
-        };
-        GreenNode {
-            kind,
-            data: GreenNodeData::Branch(branch),
+    pub fn text_len(&self) -> TextUnit {
+        match self {
+            GreenNode::Leaf(l) => l.text_len(),
+            GreenNode::Branch(b) => b.text_len(),
         }
     }
 
-    pub(crate) fn push_trivia(&mut self, kind: SyntaxKind, text: TokenText) {
-        let branch = match &mut self.data {
-            GreenNodeData::Branch(branch) => branch,
-            _ => panic!()
-        };
-        branch.text_len += TextUnit::of_str(&text);
-        let leading = &mut branch.leading_trivia;
-        branch.children.last_mut().map(|(_, t)| t).unwrap_or(leading)
-            .push(Arc::new(GreenTrivia { kind, text }));
-    }
-
-    pub(crate) fn push_child(&mut self, node: Arc<GreenNode>) {
-        let branch = match &mut self.data {
-            GreenNodeData::Branch(branch) => branch,
-            _ => panic!()
-        };
-        branch.text_len += node.text_len();
-        branch.children.push((node, Trivias::default()));
-    }
-
-    pub(crate) fn kind(&self) -> SyntaxKind {
-        self.kind
-    }
-
-    pub(crate) fn text_len(&self) -> TextUnit {
-        match &self.data {
-            GreenNodeData::Leaf(l) => l.text_len(),
-            GreenNodeData::Branch(b) => b.text_len(),
+    pub fn children(&self) -> &[GreenNode] {
+        match self {
+            GreenNode::Leaf(_) => &[],
+            GreenNode::Branch(b) => b.children(),
         }
     }
 
-    pub(crate) fn text(&self) -> String {
+    pub fn text(&self) -> String {
         let mut buff = String::new();
         go(self, &mut buff);
         return buff;
         fn go(node: &GreenNode, buff: &mut String) {
-            match &node.data {
-                GreenNodeData::Leaf(l) => buff.push_str(&l.text),
-                GreenNodeData::Branch(branch) => {
-                    add_trivia(&branch.leading_trivia, buff);
-                    branch.children.iter().for_each(|(child, trivias)| {
-                        go(child, buff);
-                        add_trivia(trivias, buff);
-                    })
+            match node {
+                GreenNode::Leaf(l) => buff.push_str(&l.text()),
+                GreenNode::Branch(b) => {
+                    b.children().iter().for_each(|child| go(child, buff))
                 }
             }
         }
-
-        fn add_trivia(trivias: &Trivias, buff: &mut String) {
-            trivias.iter().for_each(|t| buff.push_str(&t.text))
-        }
     }
+}
 
-    pub(crate) fn n_children(&self) -> usize {
-        match &self.data {
-            GreenNodeData::Leaf(_) => 0,
-            GreenNodeData::Branch(branch) => branch.children.len(),
-        }
-    }
+pub(crate) struct GreenNodeBuilder {
+    kind: SyntaxKind,
+    children: Vec<GreenNode>,
+}
 
-    pub(crate) fn nth_child(&self, idx: usize) -> &Arc<GreenNode> {
-        match &self.data {
-            GreenNodeData::Leaf(_) => panic!("leaf nodes have no children"),
-            GreenNodeData::Branch(branch) => &branch.children[idx].0,
-        }
+impl GreenNodeBuilder {
+    pub(crate) fn new_leaf(kind: SyntaxKind, text: &str) -> GreenNode {
+        GreenNode::Leaf(GreenLeaf::new(kind, text))
     }
 
-    pub(crate) fn nth_trivias(&self, idx: usize) -> &Trivias {
-        match &self.data {
-            GreenNodeData::Leaf(_) => panic!("leaf nodes have no children"),
-            GreenNodeData::Branch(branch) => if idx == 0 {
-                &branch.leading_trivia
-            } else {
-                &branch.children[idx - 1].1
-            },
+    pub(crate) fn new_internal(kind: SyntaxKind) -> GreenNodeBuilder {
+        GreenNodeBuilder {
+            kind,
+            children: Vec::new(),
         }
     }
 
-    pub(crate) fn is_leaf(&self) -> bool {
-        match self.data {
-            GreenNodeData::Leaf(_) => true,
-            GreenNodeData::Branch(_) => false
-        }
+    pub(crate) fn push_child(&mut self, node: GreenNode) {
+        self.children.push(node)
     }
-}
 
-#[derive(Debug)]
-enum GreenNodeData {
-    Leaf(GreenLeaf),
-    Branch(GreenBranch),
+    pub(crate) fn build(self) -> GreenNode {
+        let branch = GreenBranch::new(self.kind, self.children);
+        GreenNode::Branch(Arc::new(branch))
+    }
 }
 
-#[derive(Debug)]
-struct GreenLeaf {
-    text: TokenText
-}
 
-#[derive(Debug)]
-struct GreenBranch {
-    text_len: TextUnit,
-    leading_trivia: Trivias,
-    children: Vec<(Arc<GreenNode>, Trivias)>,
+#[test]
+fn assert_send_sync() {
+    fn f<T: Send + Sync>() {}
+    f::<GreenNode>();
 }
 
-#[derive(Debug)]
-pub(crate) struct GreenTrivia {
-    pub(crate) kind: SyntaxKind,
-    pub(crate) text: TokenText,
+#[derive(Clone, Debug)]
+pub(crate) enum GreenLeaf {
+    Whitespace {
+        newlines: u8,
+        spaces: u8,
+    },
+    Token {
+        kind: SyntaxKind,
+        text: Arc<str>,
+    },
 }
 
-type Trivias = Vec<Arc<GreenTrivia>>;
+impl GreenLeaf {
+    fn new(kind: SyntaxKind, text: &str) -> Self {
+        if kind == WHITESPACE {
+            let newlines = text.bytes().take_while(|&b| b == b'\n').count();
+            let spaces = text[newlines..].bytes().take_while(|&b| b == b' ').count();
+            if newlines + spaces == text.len() && newlines <= N_NEWLINES && spaces <= N_SPACES {
+                return GreenLeaf::Whitespace { newlines: newlines as u8, spaces: spaces as u8 };
+            }
+        }
+        GreenLeaf::Token { kind, text: text.to_owned().into_boxed_str().into() }
+    }
 
+    pub(crate) fn kind(&self) -> SyntaxKind {
+        match self {
+            GreenLeaf::Whitespace { .. } => WHITESPACE,
+            GreenLeaf::Token { kind, .. } => *kind,
+        }
+    }
 
-pub(crate) trait TextLen {
-    fn text_len(&self) -> TextUnit;
-}
+    pub(crate) fn text(&self) -> &str {
+        match self {
+            &GreenLeaf::Whitespace { newlines, spaces } => {
+                let newlines = newlines as usize;
+                let spaces = spaces as usize;
+                assert!(newlines <= N_NEWLINES && spaces <= N_SPACES);
+                &WS[N_NEWLINES - newlines..N_NEWLINES + spaces]
+            }
+            GreenLeaf::Token { text, .. } => text,
+        }
+    }
 
-impl TextLen for GreenTrivia {
-    fn text_len(&self) -> TextUnit {
-        TextUnit::of_str(&self.text)
+    pub(crate) fn text_len(&self) -> TextUnit {
+        TextUnit::of_str(self.text())
     }
 }
 
-impl<T: TextLen> TextLen for Arc<T> {
-    fn text_len(&self) -> TextUnit {
-        let this: &T = self;
-        this.text_len()
-    }
+const N_NEWLINES: usize = 16;
+const N_SPACES: usize = 64;
+const WS: &str =
+    "\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n                                                                ";
+
+#[derive(Clone, Debug)]
+pub(crate) struct GreenBranch {
+    text_len: TextUnit,
+    kind: SyntaxKind,
+    children: Vec<GreenNode>,
 }
 
-impl TextLen for GreenNode {
-    fn text_len(&self) -> TextUnit {
-        self.text_len()
+impl GreenBranch {
+    fn new(kind: SyntaxKind, children: Vec<GreenNode>) -> GreenBranch {
+        let text_len = children.iter().map(|x| x.text_len()).sum::<TextUnit>();
+        GreenBranch { text_len, kind, children }
     }
-}
 
-impl TextLen for GreenLeaf {
-    fn text_len(&self) -> TextUnit {
-        TextUnit::of_str(&self.text)
+    pub fn kind(&self) -> SyntaxKind {
+        self.kind
     }
-}
 
-impl TextLen for GreenBranch {
-    fn text_len(&self) -> TextUnit {
+    pub fn text_len(&self) -> TextUnit {
         self.text_len
     }
-}
 
-impl<T: TextLen> TextLen for [T] {
-    fn text_len(&self) -> TextUnit {
-        self.iter().map(TextLen::text_len).sum()
+    pub fn children(&self) -> &[GreenNode] {
+        self.children.as_slice()
     }
 }
+
diff --git a/src/yellow/mod.rs b/src/yellow/mod.rs
index 9e64d042f..9a6203cc1 100644
--- a/src/yellow/mod.rs
+++ b/src/yellow/mod.rs
@@ -8,7 +8,7 @@ use std::{
     mem
 };
 pub(crate) use self::{
-    green::{GreenNode, TextLen},
+    green::{GreenNode, GreenNodeBuilder},
     red::RedNode,
     syntax::SyntaxError,
     builder::GreenBuilder,
diff --git a/src/yellow/red.rs b/src/yellow/red.rs
index 71212a081..e2dbceeae 100644
--- a/src/yellow/red.rs
+++ b/src/yellow/red.rs
@@ -1,12 +1,12 @@
 use std::sync::{Arc, RwLock};
 use {
     TextUnit,
-    yellow::{Ptr, GreenNode, TextLen}
+    yellow::{Ptr, GreenNode},
 };
 
 #[derive(Debug)]
 pub(crate) struct RedNode {
-    green: Arc<GreenNode>,
+    green: GreenNode,
     parent: Option<ParentData>,
     children: RwLock<Vec<Option<Arc<RedNode>>>>,
 }
@@ -20,30 +20,30 @@ struct ParentData {
 
 impl RedNode {
     pub fn new_root(
-        green: Arc<GreenNode>,
+        green: GreenNode,
     ) -> RedNode {
         RedNode::new(green, None)
     }
 
     fn new_child(
-        green: Arc<GreenNode>,
+        green: GreenNode,
         parent: Ptr<RedNode>,
         start_offset: TextUnit,
-        index_in_parent: usize
+        index_in_parent: usize,
     ) -> RedNode {
         let parent_data = ParentData {
             parent,
             start_offset,
-            index_in_parent
+            index_in_parent,
         };
         RedNode::new(green, Some(parent_data))
     }
 
     fn new(
-        green: Arc<GreenNode>,
+        green: GreenNode,
         parent: Option<ParentData>,
     ) -> RedNode {
-        let children = vec![None; green.n_children()];
+        let children = vec![None; green.children().len()];
         RedNode { green, parent, children: RwLock::new(children) }
     }
 
@@ -59,29 +59,22 @@ impl RedNode {
     }
 
     pub(crate) fn n_children(&self) -> usize {
-        self.green.n_children()
+        self.green.children().len()
     }
 
-    pub(crate) fn nth_child(&self, me: Ptr<RedNode>, n: usize) -> Arc<RedNode> {
-        match &self.children.read().unwrap()[n] {
+    pub(crate) fn nth_child(&self, me: Ptr<RedNode>, idx: usize) -> Arc<RedNode> {
+        match &self.children.read().unwrap()[idx] {
             Some(child) => return child.clone(),
             None => (),
         }
         let mut children = self.children.write().unwrap();
-        if children[n].is_none() {
-            let start_offset = {
-                let mut acc = self.start_offset();
-                for i in 0..n {
-                    acc += self.green.nth_trivias(i).text_len();
-                    acc += self.green.nth_child(i).text_len();
-                }
-                acc += self.green.nth_trivias(n).text_len();
-                acc
-            };
-            let green = self.green.nth_child(n).clone();
-            let child = RedNode::new_child(green, me, start_offset, n);
-            children[n] = Some(Arc::new(child))
+        if children[idx].is_none() {
+            let green_children = self.green.children();
+            let start_offset = self.start_offset()
+                + green_children[..idx].iter().map(|x| x.text_len()).sum::<TextUnit>();
+            let child = RedNode::new_child(green_children[idx].clone(), me, start_offset, idx);
+            children[idx] = Some(Arc::new(child))
         }
-        children[n].as_ref().unwrap().clone()
+        children[idx].as_ref().unwrap().clone()
     }
 }
diff --git a/src/yellow/syntax.rs b/src/yellow/syntax.rs
index 78fa5bf95..53d8c82b9 100644
--- a/src/yellow/syntax.rs
+++ b/src/yellow/syntax.rs
@@ -6,14 +6,13 @@ use std::{
 use {
     TextRange, TextUnit,
     SyntaxKind::{self, *},
-    yellow::{Ptr, RedNode, GreenNode, TextLen},
+    yellow::{Ptr, RedNode, GreenNode},
 };
 
 #[derive(Clone)]
 pub struct SyntaxNode {
     pub(crate) root: SyntaxRoot,
     red: Ptr<RedNode>,
-    trivia_pos: Option<(usize, usize)>,
 }
 
 #[derive(Clone)]
@@ -29,80 +28,38 @@ pub(crate) struct SyntaxError {
 }
 
 impl SyntaxNode {
-    pub(crate) fn new(root: Arc<GreenNode>, errors: Vec<SyntaxError>) -> SyntaxNode {
+    pub(crate) fn new(root: GreenNode, errors: Vec<SyntaxError>) -> SyntaxNode {
         let root = Arc::new(RedNode::new_root(root));
         let red = Ptr::new(&root);
         let root = SyntaxRoot { red: root, errors: Arc::new(errors) };
-        SyntaxNode { root, red, trivia_pos: None }
+        SyntaxNode { root, red }
     }
 
     pub fn kind(&self) -> SyntaxKind {
-        let green = self.red().green();
-        match self.trivia_pos {
-            None => green.kind(),
-            Some((i, j)) => green.nth_trivias(i)[j].kind
-        }
+        self.red().green().kind()
     }
 
     pub fn range(&self) -> TextRange {
         let red = self.red();
-        let green = red.green();
-        match self.trivia_pos {
-            None => TextRange::offset_len(red.start_offset(), red.green().text_len()),
-            Some((i, j)) => {
-                let trivias = green.nth_trivias(i);
-                let offset = if i == 0 {
-                    red.start_offset()
-                } else {
-                    let prev_child = red.nth_child(Ptr::clone(&self.red), i - 1);
-                    let mut offset = prev_child.start_offset() + prev_child.green().text_len();
-                    for k in 0..j {
-                        offset += &trivias[k].text_len();
-                    }
-                    offset
-                };
-                TextRange::offset_len(offset, trivias[j].text_len())
-            }
-        }
+        TextRange::offset_len(
+            red.start_offset(),
+            red.green().text_len(),
+        )
     }
 
     pub fn text(&self) -> String {
-        let green = self.red().green();
-        match self.trivia_pos {
-            None => green.text(),
-            Some((i, j)) => green.nth_trivias(i)[j].text.clone()
-        }
+        self.red().green().text()
     }
 
     pub fn children(&self) -> Vec<SyntaxNode> {
-        let mut res = Vec::new();
         let red = self.red();
-        let green = red.green();
-        if green.is_leaf() || self.trivia_pos.is_some() {
-            return Vec::new();
-        }
-        for (j, _) in green.nth_trivias(0).iter().enumerate() {
-            res.push(SyntaxNode {
-                root: self.root.clone(),
-                red: Ptr::clone(&self.red),
-                trivia_pos: Some((0, j)),
-            })
-        }
-
         let n_children = red.n_children();
+        let mut res = Vec::with_capacity(n_children);
         for i in 0..n_children {
             res.push(SyntaxNode {
                 root: self.root.clone(),
                 red: Ptr::new(&red.nth_child(Ptr::clone(&self.red), i)),
-                trivia_pos: None,
             });
-            for (j, _) in green.nth_trivias(i + 1).iter().enumerate() {
-                res.push(SyntaxNode {
-                    root: self.root.clone(),
-                    red: self.red.clone(),
-                    trivia_pos: Some((i + 1, j)),
-                })
-            }
         }
         res
     }