1 files changed, 63 insertions, 89 deletions

diff --git a/crates/mbe/src/subtree_source.rs b/crates/mbe/src/subtree_source.rs
index d10d4b70e..d7433bd35 100644
--- a/crates/mbe/src/subtree_source.rs
+++ b/crates/mbe/src/subtree_source.rs
@@ -1,129 +1,104 @@
 //! FIXME: write short doc here
 
 use parser::{Token, TokenSource};
-use std::cell::{Cell, Ref, RefCell};
 use syntax::{lex_single_syntax_kind, SmolStr, SyntaxKind, SyntaxKind::*, T};
-use tt::buffer::{Cursor, TokenBuffer};
+use tt::buffer::TokenBuffer;
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 struct TtToken {
-    kind: SyntaxKind,
-    is_joint_to_next: bool,
+    tt: Token,
     text: SmolStr,
 }
 
-pub(crate) struct SubtreeTokenSource<'a> {
-    cached_cursor: Cell<Cursor<'a>>,
-    cached: RefCell<Vec<Option<TtToken>>>,
+pub(crate) struct SubtreeTokenSource {
+    cached: Vec<TtToken>,
     curr: (Token, usize),
 }
 
-impl<'a> SubtreeTokenSource<'a> {
+impl<'a> SubtreeTokenSource {
     // Helper function used in test
     #[cfg(test)]
     pub(crate) fn text(&self) -> SmolStr {
-        match *self.get(self.curr.1) {
+        match self.cached.get(self.curr.1) {
             Some(ref tt) => tt.text.clone(),
             _ => SmolStr::new(""),
         }
     }
 }
 
-impl<'a> SubtreeTokenSource<'a> {
-    pub(crate) fn new(buffer: &'a TokenBuffer) -> SubtreeTokenSource<'a> {
-        let cursor = buffer.begin();
+impl<'a> SubtreeTokenSource {
+    pub(crate) fn new(buffer: &TokenBuffer) -> SubtreeTokenSource {
+        let mut current = buffer.begin();
+        let mut cached = Vec::with_capacity(100);
 
-        let mut res = SubtreeTokenSource {
-            curr: (Token { kind: EOF, is_jointed_to_next: false }, 0),
-            cached_cursor: Cell::new(cursor),
-            cached: RefCell::new(Vec::with_capacity(10)),
-        };
-        res.curr = (res.mk_token(0), 0);
-        res
-    }
+        while !current.eof() {
+            let cursor = current;
+            let tt = cursor.token_tree();
 
-    fn mk_token(&self, pos: usize) -> Token {
-        match *self.get(pos) {
-            Some(ref tt) => Token { kind: tt.kind, is_jointed_to_next: tt.is_joint_to_next },
-            None => Token { kind: EOF, is_jointed_to_next: false },
-        }
-    }
-
-    fn get(&self, pos: usize) -> Ref<Option<TtToken>> {
-        fn is_lifetime(c: Cursor) -> Option<(Cursor, SmolStr)> {
-            let tkn = c.token_tree();
-
-            if let Some(tt::TokenTree::Leaf(tt::Leaf::Punct(punct))) = tkn {
+            // Check if it is lifetime
+            if let Some(tt::buffer::TokenTreeRef::Leaf(tt::Leaf::Punct(punct), _)) = tt {
                 if punct.char == '\'' {
-                    let next = c.bump();
-                    if let Some(tt::TokenTree::Leaf(tt::Leaf::Ident(ident))) = next.token_tree() {
-                        let res_cursor = next.bump();
-                        let text = SmolStr::new("'".to_string() + &ident.to_string());
-
-                        return Some((res_cursor, text));
+                    let next = cursor.bump();
+                    if let Some(tt::buffer::TokenTreeRef::Leaf(tt::Leaf::Ident(ident), _)) =
+                        next.token_tree()
+                    {
+                        let text = SmolStr::new("'".to_string() + &ident.text);
+                        cached.push(TtToken {
+                            tt: Token { kind: LIFETIME_IDENT, is_jointed_to_next: false },
+                            text,
+                        });
+                        current = next.bump();
+                        continue;
                     } else {
                         panic!("Next token must be ident : {:#?}", next.token_tree());
                     }
                 }
             }
 
-            None
-        }
-
-        if pos < self.cached.borrow().len() {
-            return Ref::map(self.cached.borrow(), |c| &c[pos]);
-        }
-
-        {
-            let mut cached = self.cached.borrow_mut();
-            while pos >= cached.len() {
-                let cursor = self.cached_cursor.get();
-                if cursor.eof() {
-                    cached.push(None);
-                    continue;
+            current = match tt {
+                Some(tt::buffer::TokenTreeRef::Leaf(leaf, _)) => {
+                    cached.push(convert_leaf(&leaf));
+                    cursor.bump()
                 }
-
-                if let Some((curr, text)) = is_lifetime(cursor) {
-                    cached.push(Some(TtToken {
-                        kind: LIFETIME_IDENT,
-                        is_joint_to_next: false,
-                        text,
-                    }));
-                    self.cached_cursor.set(curr);
-                    continue;
+                Some(tt::buffer::TokenTreeRef::Subtree(subtree, _)) => {
+                    cached.push(convert_delim(subtree.delimiter_kind(), false));
+                    cursor.subtree().unwrap()
                 }
-
-                match cursor.token_tree() {
-                    Some(tt::TokenTree::Leaf(leaf)) => {
-                        cached.push(Some(convert_leaf(&leaf)));
-                        self.cached_cursor.set(cursor.bump());
-                    }
-                    Some(tt::TokenTree::Subtree(subtree)) => {
-                        self.cached_cursor.set(cursor.subtree().unwrap());
-                        cached.push(Some(convert_delim(subtree.delimiter_kind(), false)));
-                    }
-                    None => {
-                        if let Some(subtree) = cursor.end() {
-                            cached.push(Some(convert_delim(subtree.delimiter_kind(), true)));
-                            self.cached_cursor.set(cursor.bump());
-                        }
+                None => {
+                    if let Some(subtree) = cursor.end() {
+                        cached.push(convert_delim(subtree.delimiter_kind(), true));
+                        cursor.bump()
+                    } else {
+                        continue;
                     }
                 }
-            }
+            };
         }
 
-        Ref::map(self.cached.borrow(), |c| &c[pos])
+        let mut res = SubtreeTokenSource {
+            curr: (Token { kind: EOF, is_jointed_to_next: false }, 0),
+            cached,
+        };
+        res.curr = (res.token(0), 0);
+        res
+    }
+
+    fn token(&self, pos: usize) -> Token {
+        match self.cached.get(pos) {
+            Some(it) => it.tt,
+            None => Token { kind: EOF, is_jointed_to_next: false },
+        }
     }
 }
 
-impl<'a> TokenSource for SubtreeTokenSource<'a> {
+impl<'a> TokenSource for SubtreeTokenSource {
     fn current(&self) -> Token {
         self.curr.0
     }
 
     /// Lookahead n token
     fn lookahead_nth(&self, n: usize) -> Token {
-        self.mk_token(self.curr.1 + n)
+        self.token(self.curr.1 + n)
     }
 
     /// bump cursor to next token
@@ -131,13 +106,12 @@ impl<'a> TokenSource for SubtreeTokenSource<'a> {
         if self.current().kind == EOF {
             return;
         }
-
-        self.curr = (self.mk_token(self.curr.1 + 1), self.curr.1 + 1);
+        self.curr = (self.token(self.curr.1 + 1), self.curr.1 + 1);
     }
 
     /// Is the current token a specified keyword?
     fn is_keyword(&self, kw: &str) -> bool {
-        match *self.get(self.curr.1) {
+        match self.cached.get(self.curr.1) {
             Some(ref t) => t.text == *kw,
             _ => false,
         }
@@ -155,7 +129,7 @@ fn convert_delim(d: Option<tt::DelimiterKind>, closing: bool) -> TtToken {
     let idx = closing as usize;
     let kind = kinds[idx];
     let text = if !texts.is_empty() { &texts[idx..texts.len() - (1 - idx)] } else { "" };
-    TtToken { kind, is_joint_to_next: false, text: SmolStr::new(text) }
+    TtToken { tt: Token { kind, is_jointed_to_next: false }, text: SmolStr::new(text) }
 }
 
 fn convert_literal(l: &tt::Literal) -> TtToken {
@@ -169,7 +143,7 @@ fn convert_literal(l: &tt::Literal) -> TtToken {
         })
         .unwrap_or_else(|| panic!("Fail to convert given literal {:#?}", &l));
 
-    TtToken { kind, is_joint_to_next: false, text: l.text.clone() }
+    TtToken { tt: Token { kind, is_jointed_to_next: false }, text: l.text.clone() }
 }
 
 fn convert_ident(ident: &tt::Ident) -> TtToken {
@@ -180,7 +154,7 @@ fn convert_ident(ident: &tt::Ident) -> TtToken {
         _ => SyntaxKind::from_keyword(ident.text.as_str()).unwrap_or(IDENT),
     };
 
-    TtToken { kind, is_joint_to_next: false, text: ident.text.clone() }
+    TtToken { tt: Token { kind, is_jointed_to_next: false }, text: ident.text.clone() }
 }
 
 fn convert_punct(p: tt::Punct) -> TtToken {
@@ -194,7 +168,7 @@ fn convert_punct(p: tt::Punct) -> TtToken {
         let s: &str = p.char.encode_utf8(&mut buf);
         SmolStr::new(s)
     };
-    TtToken { kind, is_joint_to_next: p.spacing == tt::Spacing::Joint, text }
+    TtToken { tt: Token { kind, is_jointed_to_next: p.spacing == tt::Spacing::Joint }, text }
 }
 
 fn convert_leaf(leaf: &tt::Leaf) -> TtToken {
@@ -208,6 +182,7 @@ fn convert_leaf(leaf: &tt::Leaf) -> TtToken {
 #[cfg(test)]
 mod tests {
     use super::{convert_literal, TtToken};
+    use parser::Token;
     use syntax::{SmolStr, SyntaxKind};
 
     #[test]
@@ -218,8 +193,7 @@ mod tests {
                 text: SmolStr::new("-42.0")
             }),
             TtToken {
-                kind: SyntaxKind::FLOAT_NUMBER,
-                is_joint_to_next: false,
+                tt: Token { kind: SyntaxKind::FLOAT_NUMBER, is_jointed_to_next: false },
                 text: SmolStr::new("-42.0")
             }
         );