3 files changed, 39 insertions, 72 deletions

diff --git a/crates/ra_mbe/src/subtree_parser.rs b/crates/ra_mbe/src/subtree_parser.rs
index 9cc989b23..1f12e42ef 100644
--- a/crates/ra_mbe/src/subtree_parser.rs
+++ b/crates/ra_mbe/src/subtree_parser.rs
@@ -2,16 +2,38 @@ use crate::subtree_source::SubtreeTokenSource;
 
 use ra_parser::{TokenSource, TreeSink};
 use ra_syntax::{SyntaxKind};
-use tt::buffer::TokenBuffer;
+use tt::buffer::{TokenBuffer, Cursor};
 
-struct OffsetTokenSink {
-    token_pos: usize,
+struct OffsetTokenSink<'a> {
+    cursor: Cursor<'a>,
     error: bool,
 }
 
-impl TreeSink for OffsetTokenSink {
+impl<'a> OffsetTokenSink<'a> {
+    pub fn collect(&self, begin: Cursor<'a>) -> Vec<tt::TokenTree> {
+        if !self.cursor.is_root() {
+            return vec![];
+        }
+
+        let mut curr = begin;
+        let mut res = vec![];
+
+        while self.cursor != curr {
+            if let Some(token) = curr.token_tree() {
+                res.push(token);
+            }
+            curr = curr.bump();
+        }
+
+        res
+    }
+}
+
+impl<'a> TreeSink for OffsetTokenSink<'a> {
     fn token(&mut self, _kind: SyntaxKind, n_tokens: u8) {
-        self.token_pos += n_tokens as usize;
+        for _ in 0..n_tokens {
+            self.cursor = self.cursor.bump_subtree();
+        }
     }
     fn start_node(&mut self, _kind: SyntaxKind) {}
     fn finish_node(&mut self) {}
@@ -72,23 +94,21 @@ impl<'a> Parser<'a> {
     {
         let buffer = TokenBuffer::new(&self.subtree.token_trees[*self.cur_pos..]);
         let mut src = SubtreeTokenSource::new(&buffer);
-        let mut sink = OffsetTokenSink { token_pos: 0, error: false };
+        let mut sink = OffsetTokenSink { cursor: buffer.begin(), error: false };
 
         f(&mut src, &mut sink);
 
-        let r = self.finish(sink.token_pos, &mut src);
+        let r = self.finish(buffer.begin(), &mut sink);
         if sink.error {
             return None;
         }
         r
     }
 
-    fn finish(self, parsed_token: usize, src: &mut SubtreeTokenSource) -> Option<tt::TokenTree> {
-        let res = src.bump_n(parsed_token);
+    fn finish(self, begin: Cursor, sink: &mut OffsetTokenSink) -> Option<tt::TokenTree> {
+        let res = sink.collect(begin);
         *self.cur_pos += res.len();
 
-        let res: Vec<_> = res.into_iter().collect();
-
         match res.len() {
             0 => None,
             1 => Some(res[0].clone()),
diff --git a/crates/ra_mbe/src/subtree_source.rs b/crates/ra_mbe/src/subtree_source.rs
index 2ba0b1601..8d6fd3e43 100644
--- a/crates/ra_mbe/src/subtree_source.rs
+++ b/crates/ra_mbe/src/subtree_source.rs
@@ -11,8 +11,7 @@ struct TtToken {
 }
 
 pub(crate) struct SubtreeTokenSource<'a> {
-    start: Cursor<'a>,
-    cursor: Cell<Cursor<'a>>,
+    cached_cursor: Cell<Cursor<'a>>,
     cached: RefCell<Vec<Option<TtToken>>>,
     curr: (Token, usize),
 }
@@ -34,19 +33,13 @@ impl<'a> SubtreeTokenSource<'a> {
 
         let mut res = SubtreeTokenSource {
             curr: (Token { kind: EOF, is_jointed_to_next: false }, 0),
-            start: cursor,
-            cursor: Cell::new(cursor),
+            cached_cursor: Cell::new(cursor),
             cached: RefCell::new(Vec::with_capacity(10)),
         };
         res.curr = (res.mk_token(0), 0);
         res
     }
 
-    pub(crate) fn bump_n(&mut self, parsed_tokens: usize) -> Vec<tt::TokenTree> {
-        let res = self.collect_token_trees(parsed_tokens);
-        res
-    }
-
     fn mk_token(&self, pos: usize) -> Token {
         match self.get(pos) {
             Some(tt) => Token { kind: tt.kind, is_jointed_to_next: tt.is_joint_to_next },
@@ -61,7 +54,7 @@ impl<'a> SubtreeTokenSource<'a> {
         }
 
         while pos >= cached.len() {
-            let cursor = self.cursor.get();
+            let cursor = self.cached_cursor.get();
             if cursor.eof() {
                 cached.push(None);
                 continue;
@@ -70,16 +63,16 @@ impl<'a> SubtreeTokenSource<'a> {
             match cursor.token_tree() {
                 Some(tt::TokenTree::Leaf(leaf)) => {
                     cached.push(Some(convert_leaf(&leaf)));
-                    self.cursor.set(cursor.bump());
+                    self.cached_cursor.set(cursor.bump());
                 }
                 Some(tt::TokenTree::Subtree(subtree)) => {
-                    self.cursor.set(cursor.subtree().unwrap());
+                    self.cached_cursor.set(cursor.subtree().unwrap());
                     cached.push(Some(convert_delim(subtree.delimiter, false)));
                 }
                 None => {
                     if let Some(subtree) = cursor.end() {
                         cached.push(Some(convert_delim(subtree.delimiter, true)));
-                        self.cursor.set(cursor.bump());
+                        self.cached_cursor.set(cursor.bump());
                     }
                 }
             }
@@ -87,48 +80,6 @@ impl<'a> SubtreeTokenSource<'a> {
 
         return cached[pos].clone();
     }
-
-    fn collect_token_trees(&self, n: usize) -> Vec<tt::TokenTree> {
-        let mut res = vec![];
-
-        let mut pos = 0;
-        let mut cursor = self.start;
-        let mut level = 0;
-
-        while pos < n {
-            if cursor.eof() {
-                break;
-            }
-
-            match cursor.token_tree() {
-                Some(tt::TokenTree::Leaf(leaf)) => {
-                    if level == 0 {
-                        res.push(leaf.into());
-                    }
-                    cursor = cursor.bump();
-                    pos += 1;
-                }
-                Some(tt::TokenTree::Subtree(subtree)) => {
-                    if level == 0 {
-                        res.push(subtree.into());
-                    }
-                    pos += 1;
-                    level += 1;
-                    cursor = cursor.subtree().unwrap();
-                }
-
-                None => {
-                    if let Some(_) = cursor.end() {
-                        level -= 1;
-                        pos += 1;
-                        cursor = cursor.bump();
-                    }
-                }
-            }
-        }
-
-        res
-    }
 }
 
 impl<'a> TokenSource for SubtreeTokenSource<'a> {
@@ -147,7 +98,7 @@ impl<'a> TokenSource for SubtreeTokenSource<'a> {
             return;
         }
 
-        self.curr = (self.mk_token(self.curr.1 + 1), self.curr.1 + 1)
+        self.curr = (self.mk_token(self.curr.1 + 1), self.curr.1 + 1);
     }
 
     /// Is the current token a specified keyword?
diff --git a/crates/ra_mbe/src/syntax_bridge.rs b/crates/ra_mbe/src/syntax_bridge.rs
index 058f65ffd..caeb842b1 100644
--- a/crates/ra_mbe/src/syntax_bridge.rs
+++ b/crates/ra_mbe/src/syntax_bridge.rs
@@ -298,11 +298,7 @@ fn delim_to_str(d: tt::Delimiter, closing: bool) -> SmolStr {
 impl<'a> TreeSink for TtTreeSink<'a> {
     fn token(&mut self, kind: SyntaxKind, n_tokens: u8) {
         if kind == L_DOLLAR || kind == R_DOLLAR {
-            if let Some(_) = self.cursor.end() {
-                self.cursor = self.cursor.bump();
-            } else {
-                self.cursor = self.cursor.subtree().unwrap();
-            }
+            self.cursor = self.cursor.bump_subtree();
             return;
         }