2 files changed, 160 insertions, 75 deletions

diff --git a/crates/ra_syntax/src/parser_impl/event.rs b/crates/ra_syntax/src/parser_impl/event.rs
index 9fd56b996..95e5ce4cc 100644
--- a/crates/ra_syntax/src/parser_impl/event.rs
+++ b/crates/ra_syntax/src/parser_impl/event.rs
@@ -9,9 +9,10 @@
 //! this stream to a real tree.
 use std::mem;
 use {
+    TextUnit, TextRange, SmolStr,
     lexer::Token,
     parser_impl::Sink,
-    SyntaxKind::{self, TOMBSTONE},
+    SyntaxKind::{self, *},
 };
 
 
@@ -78,77 +79,162 @@ pub(crate) enum Event {
     },
 }
 
+pub(super) struct EventProcessor<'a, S: Sink> {
+    sink: S,
+    text_pos: TextUnit,
+    text: &'a str,
+    token_pos: usize,
+    tokens: &'a [Token],
+    events: &'a mut [Event],
+}
 
-pub(super) fn process<'a, S: Sink<'a>>(builder: &mut S, tokens: &[Token], mut events: Vec<Event>) {
-    fn tombstone() -> Event {
-        Event::Start { kind: TOMBSTONE, forward_parent: None }
+impl<'a, S: Sink> EventProcessor<'a, S> {
+    pub(super) fn new(sink: S, text: &'a str, tokens: &'a[Token], events: &'a mut [Event]) -> EventProcessor<'a, S> {
+        EventProcessor {
+            sink,
+            text_pos: 0.into(),
+            text,
+            token_pos: 0,
+            tokens,
+            events
+        }
     }
-    let eat_ws = |idx: &mut usize, builder: &mut S| {
-        while let Some(token) = tokens.get(*idx) {
-            if !token.kind.is_trivia() {
-                break;
-            }
-            builder.leaf(token.kind, token.len);
-            *idx += 1
+
+    pub(super) fn process(mut self) -> S {
+        fn tombstone() -> Event {
+            Event::Start { kind: TOMBSTONE, forward_parent: None }
         }
-    };
-
-    let events: &mut [Event] = &mut events;
-    let mut depth = 0;
-    let mut forward_parents = Vec::new();
-    let mut next_tok_idx = 0;
-    for i in 0..events.len() {
-        match mem::replace(&mut events[i], tombstone()) {
-            Event::Start {
-                kind: TOMBSTONE, ..
-            } => (),
-
-            Event::Start { kind, forward_parent } => {
-                forward_parents.push(kind);
-                let mut idx = i;
-                let mut fp = forward_parent;
-                while let Some(fwd) = fp {
-                    idx += fwd as usize;
-                    fp = match mem::replace(&mut events[idx], tombstone()) {
-                        Event::Start {
-                            kind,
-                            forward_parent,
-                        } => {
-                            forward_parents.push(kind);
-                            forward_parent
-                        },
-                        _ => unreachable!(),
-                    };
-                }
-                for kind in forward_parents.drain(..).rev() {
-                    if depth > 0 {
-                        eat_ws(&mut next_tok_idx, builder);
+        let mut forward_parents = Vec::new();
+
+        for i in 0..self.events.len() {
+            match mem::replace(&mut self.events[i], tombstone()) {
+                Event::Start {
+                    kind: TOMBSTONE, ..
+                } => (),
+
+                Event::Start { kind, forward_parent } => {
+                    forward_parents.push(kind);
+                    let mut idx = i;
+                    let mut fp = forward_parent;
+                    while let Some(fwd) = fp {
+                        idx += fwd as usize;
+                        fp = match mem::replace(&mut self.events[idx], tombstone()) {
+                            Event::Start {
+                                kind,
+                                forward_parent,
+                            } => {
+                                forward_parents.push(kind);
+                                forward_parent
+                            },
+                            _ => unreachable!(),
+                        };
+                    }
+                    for kind in forward_parents.drain(..).rev() {
+                        self.start(kind);
                     }
-                    depth += 1;
-                    builder.start_internal(kind);
                 }
-            }
-            Event::Finish => {
-                depth -= 1;
-                if depth == 0 {
-                    eat_ws(&mut next_tok_idx, builder);
+                Event::Finish => {
+                    let last = i == self.events.len() - 1;
+                    self.finish(last);
+                },
+                Event::Token { kind, n_raw_tokens } => {
+                    self.eat_ws();
+                    let n_raw_tokens = n_raw_tokens as usize;
+                    let len = self.tokens[self.token_pos..self.token_pos + n_raw_tokens]
+                        .iter()
+                        .map(|it| it.len)
+                        .sum::<TextUnit>();
+                    self.leaf(kind, len, n_raw_tokens);
                 }
+                Event::Error { msg } => self.sink.error(msg, self.text_pos),
+            }
+        }
+        self.sink
+    }
+
+    fn start(&mut self, kind: SyntaxKind) {
+        if kind == ROOT {
+            self.sink.start_internal(kind);
+            return;
+        }
+        let n_trivias = self.tokens[self.token_pos..]
+            .iter()
+            .take_while(|it| it.kind.is_trivia())
+            .count();
+        let leading_trivias = &self.tokens[self.token_pos..self.token_pos + n_trivias];
+        let mut trivia_end = self.text_pos + leading_trivias
+            .iter()
+            .map(|it| it.len)
+            .sum::<TextUnit>();
+
+        let n_attached_trivias = {
+            let leading_trivias = leading_trivias.iter().rev()
+                .map(|it| {
+                    let next_end = trivia_end - it.len;
+                    let range = TextRange::from_to(next_end, trivia_end);
+                    trivia_end = next_end;
+                    (it.kind, &self.text[range])
+                });
+            n_attached_trivias(kind, leading_trivias)
+        };
+        self.eat_n_trivias(n_trivias - n_attached_trivias);
+        self.sink.start_internal(kind);
+        self.eat_n_trivias(n_attached_trivias);
+    }
 
-                builder.finish_internal();
+    fn finish(&mut self, last: bool) {
+        if last {
+            self.eat_ws()
+        }
+        self.sink.finish_internal();
+    }
+
+    fn eat_ws(&mut self) {
+        while let Some(&token) = self.tokens.get(self.token_pos) {
+            if !token.kind.is_trivia() {
+                break;
             }
-            Event::Token {
-                kind,
-                mut n_raw_tokens,
-            } => {
-                eat_ws(&mut next_tok_idx, builder);
-                let mut len = 0.into();
-                for _ in 0..n_raw_tokens {
-                    len += tokens[next_tok_idx].len;
-                    next_tok_idx += 1;
+            self.leaf(token.kind, token.len, 1);
+        }
+    }
+
+    fn eat_n_trivias(&mut self, n: usize) {
+        for _ in 0..n {
+            let token = self.tokens[self.token_pos];
+            assert!(token.kind.is_trivia());
+            self.leaf(token.kind, token.len, 1);
+        }
+    }
+
+    fn leaf(&mut self, kind: SyntaxKind, len: TextUnit, n_tokens: usize) {
+        let range = TextRange::offset_len(self.text_pos, len);
+        let text: SmolStr = self.text[range].into();
+        self.text_pos += len;
+        self.token_pos += n_tokens;
+        self.sink.leaf(kind, text);
+    }
+}
+
+fn n_attached_trivias<'a>(kind: SyntaxKind, trivias: impl Iterator<Item=(SyntaxKind, &'a str)>) -> usize {
+    match kind {
+        STRUCT_DEF | ENUM_DEF | FN_DEF | TRAIT_DEF | MODULE => {
+            let mut res = 0;
+            for (i, (kind, text)) in trivias.enumerate() {
+                match kind {
+                    WHITESPACE => {
+                        if text.contains("\n\n") {
+                            break;
+                        }
+                    }
+                    COMMENT => {
+                        res = i + 1;
+                    }
+                    _ => (),
                 }
-                builder.leaf(kind, len);
             }
-            Event::Error { msg } => builder.error(msg),
+            res
         }
+        _ => 0,
     }
+
 }
diff --git a/crates/ra_syntax/src/parser_impl/mod.rs b/crates/ra_syntax/src/parser_impl/mod.rs
index b343b404f..8d74cef0e 100644
--- a/crates/ra_syntax/src/parser_impl/mod.rs
+++ b/crates/ra_syntax/src/parser_impl/mod.rs
@@ -4,45 +4,44 @@ mod input;
 use std::cell::Cell;
 
 use {
+    TextUnit, SmolStr,
     lexer::Token,
     parser_api::Parser,
     parser_impl::{
-        event::{process, Event},
+        event::{EventProcessor, Event},
         input::{InputPosition, ParserInput},
     },
-    TextUnit,
 };
 
 use SyntaxKind::{self, EOF, TOMBSTONE};
 
-pub(crate) trait Sink<'a> {
+pub(crate) trait Sink {
     type Tree;
 
-    fn new(text: &'a str) -> Self;
-
-    fn leaf(&mut self, kind: SyntaxKind, len: TextUnit);
+    fn leaf(&mut self, kind: SyntaxKind, text: SmolStr);
     fn start_internal(&mut self, kind: SyntaxKind);
     fn finish_internal(&mut self);
-    fn error(&mut self, err: String);
+    fn error(&mut self, message: String, offset: TextUnit);
     fn finish(self) -> Self::Tree;
 }
 
 /// Parse a sequence of tokens into the representative node tree
-pub(crate) fn parse_with<'a, S: Sink<'a>>(
-    text: &'a str,
+pub(crate) fn parse_with<S: Sink>(
+    sink: S,
+    text: &str,
     tokens: &[Token],
     parser: fn(&mut Parser),
 ) -> S::Tree {
-    let events = {
+    let mut events = {
         let input = input::ParserInput::new(text, tokens);
         let parser_impl = ParserImpl::new(&input);
         let mut parser_api = Parser(parser_impl);
         parser(&mut parser_api);
         parser_api.0.into_events()
     };
-    let mut sink = S::new(text);
-    process(&mut sink, tokens, events);
-    sink.finish()
+    EventProcessor::new(sink, text, tokens, &mut events)
+        .process()
+        .finish()
 }
 
 /// Implementation details of `Parser`, extracted