Use SubtreeWalker instread of flatten TtToken

4 files changed, 421 insertions, 258 deletions

diff --git a/crates/ra_mbe/src/lib.rs b/crates/ra_mbe/src/lib.rs
index 84ce2b783..a21ea4dbc 100644
--- a/crates/ra_mbe/src/lib.rs
+++ b/crates/ra_mbe/src/lib.rs
@@ -383,8 +383,22 @@ SOURCE_FILE@[0; 40)
         assert_eq!(to_literal(&stm_tokens[15 + 3]).text, "\"rust1\"");
     }
 
-    /// The following tests are port from intellij-rust directly
-    /// https://github.com/intellij-rust/intellij-rust/blob/c4e9feee4ad46e7953b1948c112533360b6087bb/src/test/kotlin/org/rust/lang/core/macros/RsMacroExpansionTest.kt
+    #[test]
+    fn test_two_idents() {
+        let rules = create_rules(
+            r#"
+        macro_rules! foo {
+            ($ i:ident, $ j:ident) => {
+                fn foo() { let a = $ i; let b = $j; }
+            }
+        }
+"#,
+        );
+        assert_expansion(&rules, "foo! { foo, bar }", "fn foo () {let a = foo ; let b = bar ;}");
+    }
+
+    // The following tests are port from intellij-rust directly
+    // https://github.com/intellij-rust/intellij-rust/blob/c4e9feee4ad46e7953b1948c112533360b6087bb/src/test/kotlin/org/rust/lang/core/macros/RsMacroExpansionTest.kt
 
     #[test]
     fn test_path() {
@@ -401,7 +415,21 @@ SOURCE_FILE@[0; 40)
         assert_expansion(
             &rules,
             "foo! { bar::<u8>::baz::<u8> }",
-            "fn foo () {let a = bar :: < u8 > :: baz :: < u8 > ;}",
+            "fn foo () {let a = bar ::< u8 > ::baz ::< u8 > ;}",
+        );
+    }
+
+    #[test]
+    fn test_two_paths() {
+        let rules = create_rules(
+            r#"
+        macro_rules! foo {
+            ($ i:path, $ j:path) => {
+                fn foo() { let a = $ i; let b = $j; }
+            }
+        }
+"#,
         );
+        assert_expansion(&rules, "foo! { foo, bar }", "fn foo () {let a = foo ; let b = bar ;}");
     }
 }
diff --git a/crates/ra_mbe/src/subtree_parser.rs b/crates/ra_mbe/src/subtree_parser.rs
index 48eee6fa7..f198c8224 100644
--- a/crates/ra_mbe/src/subtree_parser.rs
+++ b/crates/ra_mbe/src/subtree_parser.rs
@@ -18,12 +18,12 @@ impl TreeSink for OffsetTokenSink {
 
 pub(crate) struct Parser<'a> {
     subtree: &'a tt::Subtree,
-    pos: &'a mut usize,
+    cur_pos: &'a mut usize,
 }
 
 impl<'a> Parser<'a> {
-    pub fn new(pos: &'a mut usize, subtree: &'a tt::Subtree) -> Parser<'a> {
-        Parser { pos, subtree }
+    pub fn new(cur_pos: &'a mut usize, subtree: &'a tt::Subtree) -> Parser<'a> {
+        Parser { cur_pos, subtree }
     }
 
     pub fn parse_path(self) -> Option<tt::TokenTree> {
@@ -35,7 +35,7 @@ impl<'a> Parser<'a> {
         F: FnOnce(&dyn TokenSource, &mut dyn TreeSink),
     {
         let mut src = SubtreeTokenSource::new(self.subtree);
-        src.advance(*self.pos, true);
+        src.start_from_nth(*self.cur_pos);
         let mut sink = OffsetTokenSink { token_pos: 0 };
 
         f(&src, &mut sink);
@@ -44,7 +44,7 @@ impl<'a> Parser<'a> {
     }
 
     fn finish(self, parsed_token: usize, src: &mut SubtreeTokenSource) -> Option<tt::TokenTree> {
-        let res = src.bump_n(parsed_token, self.pos);
+        let res = src.bump_n(parsed_token, self.cur_pos);
         let res: Vec<_> = res.into_iter().cloned().collect();
 
         match res.len() {
diff --git a/crates/ra_mbe/src/subtree_source.rs b/crates/ra_mbe/src/subtree_source.rs
index d9ba5d3d0..9dd475f2c 100644
--- a/crates/ra_mbe/src/subtree_source.rs
+++ b/crates/ra_mbe/src/subtree_source.rs
@@ -1,7 +1,8 @@
 use ra_parser::{TokenSource};
 use ra_syntax::{classify_literal, SmolStr, SyntaxKind, SyntaxKind::*};
+use std::cell::{RefCell};
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Eq, PartialEq)]
 struct TtToken {
     pub kind: SyntaxKind,
     pub is_joint_to_next: bool,
@@ -9,107 +10,319 @@ struct TtToken {
     pub n_tokens: usize,
 }
 
-/// Querier let outside to query internal tokens as string
-pub(crate) struct Querier<'a> {
-    src: &'a SubtreeTokenSource<'a>,
+#[derive(Debug, Clone, Eq, PartialEq)]
+enum WalkIndex {
+    DelimiterBegin(Option<TtToken>),
+    Token(usize, Option<TtToken>),
+    DelimiterEnd(Option<TtToken>),
+    Eof,
 }
 
-impl<'a> Querier<'a> {
-    pub(crate) fn token(&self, uidx: usize) -> (SyntaxKind, &SmolStr) {
-        let tkn = &self.src.tokens[uidx];
-        (tkn.kind, &tkn.text)
+impl<'a> SubTreeWalker<'a> {
+    fn new(subtree: &tt::Subtree) -> SubTreeWalker {
+        let mut res = SubTreeWalker {
+            pos: 0,
+            stack: vec![],
+            idx: WalkIndex::Eof,
+            last_steps: vec![],
+            subtree,
+        };
+
+        res.reset();
+        res
+    }
+
+    fn reset(&mut self) {
+        self.pos = 0;
+        self.stack = vec![(self.subtree, None)];
+        self.idx = WalkIndex::DelimiterBegin(convert_delim(self.subtree.delimiter, false));
+        self.last_steps = vec![];
+
+        while self.is_empty_delimiter() {
+            self.forward_unchecked();
+        }
+    }
+
+    // This funciton will fast forward the pos cursor,
+    // Such that backward will stop at `start_pos` point
+    fn start_from_nth(&mut self, start_pos: usize) {
+        self.reset();
+        self.pos = start_pos;
+        self.idx = self.walk_token(start_pos, false);
+
+        while self.is_empty_delimiter() {
+            self.forward_unchecked();
+        }
+    }
+
+    fn current(&self) -> Option<&TtToken> {
+        match &self.idx {
+            WalkIndex::DelimiterBegin(t) => t.as_ref(),
+            WalkIndex::Token(_, t) => t.as_ref(),
+            WalkIndex::DelimiterEnd(t) => t.as_ref(),
+            WalkIndex::Eof => None,
+        }
+    }
+
+    fn is_empty_delimiter(&self) -> bool {
+        match &self.idx {
+            WalkIndex::DelimiterBegin(None) => true,
+            WalkIndex::DelimiterEnd(None) => true,
+            _ => false,
+        }
+    }
+
+    fn backward(&mut self) {
+        if self.last_steps.is_empty() {
+            return;
+        }
+        self.pos -= 1;
+        loop {
+            self.backward_unchecked();
+            // Skip Empty delimiter
+            if self.last_steps.is_empty() || !self.is_empty_delimiter() {
+                break;
+            }
+        }
+    }
+
+    fn backward_unchecked(&mut self) {
+        if self.last_steps.is_empty() {
+            return;
+        }
+
+        let last_step = self.last_steps.pop().unwrap();
+        let do_walk_token = match self.idx {
+            WalkIndex::DelimiterBegin(_) => None,
+            WalkIndex::Token(u, _) => Some(u),
+            WalkIndex::DelimiterEnd(_) => {
+                let (top, _) = self.stack.last().unwrap();
+                Some(top.token_trees.len())
+            }
+            WalkIndex::Eof => None,
+        };
+
+        self.idx = match do_walk_token {
+            Some(u) if last_step > u => WalkIndex::DelimiterBegin(convert_delim(
+                self.stack.last().unwrap().0.delimiter,
+                false,
+            )),
+            Some(u) => self.walk_token(u - last_step, true),
+            None => match self.idx {
+                WalkIndex::Eof => {
+                    self.stack.push((self.subtree, None));
+                    WalkIndex::DelimiterEnd(convert_delim(
+                        self.stack.last().unwrap().0.delimiter,
+                        true,
+                    ))
+                }
+                _ => {
+                    let (_, last_top_idx) = self.stack.pop().unwrap();
+                    assert!(!self.stack.is_empty());
+
+                    match last_top_idx.unwrap() {
+                        0 => WalkIndex::DelimiterBegin(convert_delim(
+                            self.stack.last().unwrap().0.delimiter,
+                            false,
+                        )),
+                        c => self.walk_token(c - 1, true),
+                    }
+                }
+            },
+        };
+    }
+
+    fn forward(&mut self) {
+        self.pos += 1;
+        loop {
+            self.forward_unchecked();
+            if !self.is_empty_delimiter() {
+                break;
+            }
+        }
+    }
+
+    fn forward_unchecked(&mut self) {
+        if self.idx == WalkIndex::Eof {
+            return;
+        }
+
+        let step = self.current().map(|x| x.n_tokens).unwrap_or(1);
+        self.last_steps.push(step);
+
+        let do_walk_token = match self.idx {
+            WalkIndex::DelimiterBegin(_) => Some(0),
+            WalkIndex::Token(u, _) => Some(u + step),
+            WalkIndex::DelimiterEnd(_) => None,
+            _ => unreachable!(),
+        };
+
+        let (top, _) = self.stack.last().unwrap();
+
+        self.idx = match do_walk_token {
+            Some(u) if u >= top.token_trees.len() => {
+                WalkIndex::DelimiterEnd(convert_delim(self.stack.last().unwrap().0.delimiter, true))
+            }
+            Some(u) => self.walk_token(u, false),
+            None => {
+                let (_, last_top_idx) = self.stack.pop().unwrap();
+                match self.stack.last() {
+                    Some(top) => match last_top_idx.unwrap() {
+                        idx if idx + 1 >= top.0.token_trees.len() => {
+                            WalkIndex::DelimiterEnd(convert_delim(top.0.delimiter, true))
+                        }
+                        idx => self.walk_token(idx + 1, false),
+                    },
+
+                    None => WalkIndex::Eof,
+                }
+            }
+        };
+    }
+
+    fn walk_token(&mut self, pos: usize, backward: bool) -> WalkIndex {
+        let (top, _) = self.stack.last().unwrap();
+        match &top.token_trees[pos] {
+            tt::TokenTree::Subtree(subtree) => {
+                self.stack.push((subtree, Some(pos)));
+                let delim = convert_delim(self.stack.last().unwrap().0.delimiter, backward);
+                if backward {
+                    WalkIndex::DelimiterEnd(delim)
+                } else {
+                    WalkIndex::DelimiterBegin(delim)
+                }
+            }
+            tt::TokenTree::Leaf(leaf) => WalkIndex::Token(pos, Some(self.walk_leaf(leaf, pos))),
+        }
+    }
+
+    fn walk_leaf(&mut self, leaf: &tt::Leaf, pos: usize) -> TtToken {
+        match leaf {
+            tt::Leaf::Literal(l) => convert_literal(l),
+            tt::Leaf::Ident(ident) => convert_ident(ident),
+            tt::Leaf::Punct(punct) => {
+                let (top, _) = self.stack.last().unwrap();
+                convert_punct(punct, top, pos)
+            }
+        }
     }
 }
 
-pub(crate) struct SubtreeTokenSource<'a> {
-    tt_pos: usize,
-    tokens: Vec<TtToken>,
-    subtree: &'a tt::Subtree,
+pub(crate) trait Querier {
+    fn token(&self, uidx: usize) -> (SyntaxKind, SmolStr);
 }
 
-impl<'a> SubtreeTokenSource<'a> {
-    pub fn new(subtree: &tt::Subtree) -> SubtreeTokenSource {
-        SubtreeTokenSource { tokens: TtTokenBuilder::build(subtree), tt_pos: 0, subtree }
+// A wrapper class for ref cell
+pub(crate) struct WalkerOwner<'a> {
+    walker: RefCell<SubTreeWalker<'a>>,
+    offset: usize,
+}
+
+impl<'a> WalkerOwner<'a> {
+    fn token_idx<'b>(&self, pos: usize) -> Option<TtToken> {
+        self.set_walker_pos(pos);
+        self.walker.borrow().current().cloned()
     }
 
-    // Advance token source and skip the first delimiter
-    pub fn advance(&mut self, n_token: usize, skip_first_delimiter: bool) {
-        if skip_first_delimiter {
-            self.tt_pos += 1;
-        }
+    fn start_from_nth(&mut self, pos: usize) {
+        self.offset = pos;
+        self.walker.borrow_mut().start_from_nth(pos);
+    }
 
-        // Matching `TtToken` cursor to `tt::TokenTree` cursor
-        // It is because TtToken is not One to One mapping to tt::Token
-        // There are 3 case (`TtToken` <=> `tt::TokenTree`) :
-        // * One to One =>  ident, single char punch
-        // * Many to One => `tt::TokenTree::SubTree`
-        // * One to Many => multibyte punct
-        //
-        // Such that we cannot simpliy advance the cursor
-        // We have to bump it one by one
-        let mut pos = 0;
-        while pos < n_token {
-            pos += self.bump(&self.subtree.token_trees[pos]);
+    fn set_walker_pos(&self, mut pos: usize) {
+        pos += self.offset;
+        let mut walker = self.walker.borrow_mut();
+        while pos > walker.pos {
+            walker.forward();
+        }
+        while pos < walker.pos {
+            walker.backward();
         }
+        assert!(pos == walker.pos);
     }
 
-    pub fn querier(&self) -> Querier {
-        Querier { src: self }
+    fn new(subtree: &'a tt::Subtree) -> Self {
+        WalkerOwner { walker: RefCell::new(SubTreeWalker::new(subtree)), offset: 0 }
     }
 
-    pub(crate) fn bump_n(
-        &mut self,
-        n_tt_tokens: usize,
-        token_pos: &mut usize,
-    ) -> Vec<&tt::TokenTree> {
+    fn collect_token_tree(&mut self, n: usize) -> Vec<&tt::TokenTree> {
+        self.start_from_nth(self.offset);
+
         let mut res = vec![];
-        // Matching `TtToken` cursor to `tt::TokenTree` cursor
-        // It is because TtToken is not One to One mapping to tt::Token
-        // There are 3 case (`TtToken` <=> `tt::TokenTree`) :
-        // * One to One =>  ident, single char punch
-        // * Many to One => `tt::TokenTree::SubTree`
-        // * One to Many => multibyte punct
-        //
-        // Such that we cannot simpliy advance the cursor
-        // We have to bump it one by one
-        let next_pos = self.tt_pos + n_tt_tokens;
-
-        while self.tt_pos < next_pos {
-            let current = &self.subtree.token_trees[*token_pos];
-            let n = self.bump(current);
-            res.extend((0..n).map(|i| &self.subtree.token_trees[*token_pos + i]));
-            *token_pos += n;
+        let mut walker = self.walker.borrow_mut();
+
+        while walker.pos - self.offset < n {
+            if let WalkIndex::Token(u, tt) = &walker.idx {
+                if walker.stack.len() == 1 {
+                    // We only collect the topmost child
+                    res.push(&walker.stack[0].0.token_trees[*u]);
+                    if let Some(tt) = tt {
+                        for i in 0..tt.n_tokens - 1 {
+                            res.push(&walker.stack[0].0.token_trees[u + i]);
+                        }
+                    }
+                }
+            }
+
+            walker.forward();
         }
 
         res
     }
+}
+
+impl<'a> Querier for WalkerOwner<'a> {
+    fn token(&self, uidx: usize) -> (SyntaxKind, SmolStr) {
+        let tkn = self.token_idx(uidx).unwrap();
+        (tkn.kind, tkn.text)
+    }
+}
+
+pub(crate) struct SubtreeTokenSource<'a> {
+    walker: WalkerOwner<'a>,
+}
+
+impl<'a> SubtreeTokenSource<'a> {
+    pub fn new(subtree: &tt::Subtree) -> SubtreeTokenSource {
+        SubtreeTokenSource { walker: WalkerOwner::new(subtree) }
+    }
+
+    pub fn start_from_nth(&mut self, n: usize) {
+        self.walker.start_from_nth(n);
+    }
 
-    fn count(&self, tt: &tt::TokenTree) -> usize {
-        assert!(!self.tokens.is_empty());
-        TtTokenBuilder::count_tt_tokens(tt, None)
+    pub fn querier<'b>(&'a self) -> &'b WalkerOwner<'a>
+    where
+        'a: 'b,
+    {
+        &self.walker
     }
 
-    fn bump(&mut self, tt: &tt::TokenTree) -> usize {
-        let cur = &self.tokens[self.tt_pos];
-        let n_tokens = cur.n_tokens;
-        self.tt_pos += self.count(tt);
-        n_tokens
+    pub(crate) fn bump_n(
+        &mut self,
+        parsed_tokens: usize,
+        cursor_pos: &mut usize,
+    ) -> Vec<&tt::TokenTree> {
+        let res = self.walker.collect_token_tree(parsed_tokens);
+        *cursor_pos += res.len();
+
+        res
     }
 }
 
 impl<'a> TokenSource for SubtreeTokenSource<'a> {
     fn token_kind(&self, pos: usize) -> SyntaxKind {
-        if let Some(tok) = self.tokens.get(self.tt_pos + pos) {
+        if let Some(tok) = self.walker.token_idx(pos) {
             tok.kind
         } else {
             SyntaxKind::EOF
         }
     }
     fn is_token_joint_to_next(&self, pos: usize) -> bool {
-        self.tokens[self.tt_pos + pos].is_joint_to_next
+        self.walker.token_idx(pos).unwrap().is_joint_to_next
     }
     fn is_keyword(&self, pos: usize, kw: &str) -> bool {
-        self.tokens[self.tt_pos + pos].text == *kw
+        self.walker.token_idx(pos).unwrap().text == *kw
     }
 }
 
@@ -136,10 +349,6 @@ where
         TokenPeek { iter: itertools::multipeek(iter) }
     }
 
-    pub fn next(&mut self) -> Option<&tt::TokenTree> {
-        self.iter.next()
-    }
-
     fn current_punct2(&mut self, p: &tt::Punct) -> Option<((char, char), bool)> {
         if p.spacing != tt::Spacing::Joint {
             return None;
@@ -162,191 +371,117 @@ where
     }
 }
 
-struct TtTokenBuilder {
-    tokens: Vec<TtToken>,
-}
-
-impl TtTokenBuilder {
-    fn build(sub: &tt::Subtree) -> Vec<TtToken> {
-        let mut res = TtTokenBuilder { tokens: vec![] };
-        res.convert_subtree(sub);
-        res.tokens
-    }
-
-    fn convert_subtree(&mut self, sub: &tt::Subtree) {
-        self.push_delim(sub.delimiter, false);
-        let mut peek = TokenPeek::new(sub.token_trees.iter());
-        while let Some(tt) = peek.iter.next() {
-            self.convert_tt(tt, &mut peek);
+fn convert_multi_char_punct<'b, I>(
+    p: &tt::Punct,
+    iter: &mut TokenPeek<'b, I>,
+) -> Option<(SyntaxKind, bool, &'static str, usize)>
+where
+    I: Iterator<Item = &'b tt::TokenTree>,
+{
+    if let Some((m, is_joint_to_next)) = iter.current_punct3(p) {
+        if let Some((kind, text)) = match m {
+            ('<', '<', '=') => Some((SHLEQ, "<<=")),
+            ('>', '>', '=') => Some((SHREQ, ">>=")),
+            ('.', '.', '.') => Some((DOTDOTDOT, "...")),
+            ('.', '.', '=') => Some((DOTDOTEQ, "..=")),
+            _ => None,
+        } {
+            return Some((kind, is_joint_to_next, text, 3));
         }
-        self.push_delim(sub.delimiter, true)
     }
 
-    fn convert_tt<'b, I>(&mut self, tt: &tt::TokenTree, iter: &mut TokenPeek<'b, I>)
-    where
-        I: Iterator<Item = &'b tt::TokenTree>,
-    {
-        match tt {
-            tt::TokenTree::Leaf(token) => self.convert_token(token, iter),
-            tt::TokenTree::Subtree(sub) => self.convert_subtree(sub),
+    if let Some((m, is_joint_to_next)) = iter.current_punct2(p) {
+        if let Some((kind, text)) = match m {
+            ('<', '<') => Some((SHL, "<<")),
+            ('>', '>') => Some((SHR, ">>")),
+
+            ('|', '|') => Some((PIPEPIPE, "||")),
+            ('&', '&') => Some((AMPAMP, "&&")),
+            ('%', '=') => Some((PERCENTEQ, "%=")),
+            ('*', '=') => Some((STAREQ, "*=")),
+            ('/', '=') => Some((SLASHEQ, "/=")),
+            ('^', '=') => Some((CARETEQ, "^=")),
+
+            ('&', '=') => Some((AMPEQ, "&=")),
+            ('|', '=') => Some((PIPEEQ, "|=")),
+            ('-', '=') => Some((MINUSEQ, "-=")),
+            ('+', '=') => Some((PLUSEQ, "+=")),
+            ('>', '=') => Some((GTEQ, ">=")),
+            ('<', '=') => Some((LTEQ, "<=")),
+
+            ('-', '>') => Some((THIN_ARROW, "->")),
+            ('!', '=') => Some((NEQ, "!=")),
+            ('=', '>') => Some((FAT_ARROW, "=>")),
+            ('=', '=') => Some((EQEQ, "==")),
+            ('.', '.') => Some((DOTDOT, "..")),
+            (':', ':') => Some((COLONCOLON, "::")),
+
+            _ => None,
+        } {
+            return Some((kind, is_joint_to_next, text, 2));
         }
     }
 
-    fn convert_token<'b, I>(&mut self, token: &tt::Leaf, iter: &mut TokenPeek<'b, I>)
-    where
-        I: Iterator<Item = &'b tt::TokenTree>,
-    {
-        let tok = match token {
-            tt::Leaf::Literal(l) => TtToken {
-                kind: classify_literal(&l.text).unwrap().kind,
-                is_joint_to_next: false,
-                text: l.text.clone(),
-                n_tokens: 1,
-            },
-            tt::Leaf::Punct(p) => {
-                if let Some((kind, is_joint_to_next, text, size)) =
-                    Self::convert_multi_char_punct(p, iter)
-                {
-                    for _ in 0..size - 1 {
-                        iter.next();
-                    }
-
-                    TtToken { kind, is_joint_to_next, text: text.into(), n_tokens: size }
-                } else {
-                    let kind = match p.char {
-                        // lexer may produce combpund tokens for these ones
-                        '.' => DOT,
-                        ':' => COLON,
-                        '=' => EQ,
-                        '!' => EXCL,
-                        '-' => MINUS,
-                        c => SyntaxKind::from_char(c).unwrap(),
-                    };
-                    let text = {
-                        let mut buf = [0u8; 4];
-                        let s: &str = p.char.encode_utf8(&mut buf);
-                        SmolStr::new(s)
-                    };
-                    TtToken {
-                        kind,
-                        is_joint_to_next: p.spacing == tt::Spacing::Joint,
-                        text,
-                        n_tokens: 1,
-                    }
-                }
-            }
-            tt::Leaf::Ident(ident) => {
-                let kind = SyntaxKind::from_keyword(ident.text.as_str()).unwrap_or(IDENT);
-                TtToken { kind, is_joint_to_next: false, text: ident.text.clone(), n_tokens: 1 }
-            }
-        };
-        self.tokens.push(tok)
-    }
-
-    fn convert_multi_char_punct<'b, I>(
-        p: &tt::Punct,
-        iter: &mut TokenPeek<'b, I>,
-    ) -> Option<(SyntaxKind, bool, &'static str, usize)>
-    where
-        I: Iterator<Item = &'b tt::TokenTree>,
-    {
-        if let Some((m, is_joint_to_next)) = iter.current_punct3(p) {
-            if let Some((kind, text)) = match m {
-                ('<', '<', '=') => Some((SHLEQ, "<<=")),
-                ('>', '>', '=') => Some((SHREQ, ">>=")),
-                ('.', '.', '.') => Some((DOTDOTDOT, "...")),
-                ('.', '.', '=') => Some((DOTDOTEQ, "..=")),
-                _ => None,
-            } {
-                return Some((kind, is_joint_to_next, text, 3));
-            }
-        }
+    None
+}
 
-        if let Some((m, is_joint_to_next)) = iter.current_punct2(p) {
-            if let Some((kind, text)) = match m {
-                ('<', '<') => Some((SHL, "<<")),
-                ('>', '>') => Some((SHR, ">>")),
-
-                ('|', '|') => Some((PIPEPIPE, "||")),
-                ('&', '&') => Some((AMPAMP, "&&")),
-                ('%', '=') => Some((PERCENTEQ, "%=")),
-                ('*', '=') => Some((STAREQ, "*=")),
-                ('/', '=') => Some((SLASHEQ, "/=")),
-                ('^', '=') => Some((CARETEQ, "^=")),
-
-                ('&', '=') => Some((AMPEQ, "&=")),
-                ('|', '=') => Some((PIPEEQ, "|=")),
-                ('-', '=') => Some((MINUSEQ, "-=")),
-                ('+', '=') => Some((PLUSEQ, "+=")),
-                ('>', '=') => Some((GTEQ, ">=")),
-                ('<', '=') => Some((LTEQ, "<=")),
-
-                ('-', '>') => Some((THIN_ARROW, "->")),
-                ('!', '=') => Some((NEQ, "!=")),
-                ('=', '>') => Some((FAT_ARROW, "=>")),
-                ('=', '=') => Some((EQEQ, "==")),
-                ('.', '.') => Some((DOTDOT, "..")),
-                (':', ':') => Some((COLONCOLON, "::")),
-
-                _ => None,
-            } {
-                return Some((kind, is_joint_to_next, text, 2));
-            }
-        }
+struct SubTreeWalker<'a> {
+    pos: usize,
+    stack: Vec<(&'a tt::Subtree, Option<usize>)>,
+    idx: WalkIndex,
+    last_steps: Vec<usize>,
+    subtree: &'a tt::Subtree,
+}
 
-        None
-    }
+fn convert_delim(d: tt::Delimiter, closing: bool) -> Option<TtToken> {
+    let (kinds, texts) = match d {
+        tt::Delimiter::Parenthesis => ([L_PAREN, R_PAREN], "()"),
+        tt::Delimiter::Brace => ([L_CURLY, R_CURLY], "{}"),
+        tt::Delimiter::Bracket => ([L_BRACK, R_BRACK], "[]"),
+        tt::Delimiter::None => return None,
+    };
+
+    let idx = closing as usize;
+    let kind = kinds[idx];
+    let text = &texts[idx..texts.len() - (1 - idx)];
+    Some(TtToken { kind, is_joint_to_next: false, text: SmolStr::new(text), n_tokens: 1 })
+}
 
-    fn push_delim(&mut self, d: tt::Delimiter, closing: bool) {
-        let (kinds, texts) = match d {
-            tt::Delimiter::Parenthesis => ([L_PAREN, R_PAREN], "()"),
-            tt::Delimiter::Brace => ([L_CURLY, R_CURLY], "{}"),
-            tt::Delimiter::Bracket => ([L_BRACK, R_BRACK], "[]"),
-            tt::Delimiter::None => return,
-        };
-        let idx = closing as usize;
-        let kind = kinds[idx];
-        let text = &texts[idx..texts.len() - (1 - idx)];
-        let tok = TtToken { kind, is_joint_to_next: false, text: SmolStr::new(text), n_tokens: 1 };
-        self.tokens.push(tok)
-    }
-
-    fn skip_sibling_leaf(leaf: &tt::Leaf, iter: &mut std::slice::Iter<tt::TokenTree>) {
-        if let tt::Leaf::Punct(p) = leaf {
-            let mut peek = TokenPeek::new(iter);
-            if let Some((_, _, _, size)) = TtTokenBuilder::convert_multi_char_punct(p, &mut peek) {
-                for _ in 0..size - 1 {
-                    peek.next();
-                }
-            }
-        }
+fn convert_literal(l: &tt::Literal) -> TtToken {
+    TtToken {
+        kind: classify_literal(&l.text).unwrap().kind,
+        is_joint_to_next: false,
+        text: l.text.clone(),
+        n_tokens: 1,
     }
+}
 
-    fn count_tt_tokens(
-        tt: &tt::TokenTree,
-        iter: Option<&mut std::slice::Iter<tt::TokenTree>>,
-    ) -> usize {
-        match tt {
-            tt::TokenTree::Subtree(sub_tree) => {
-                let mut iter = sub_tree.token_trees.iter();
-                let mut count = match sub_tree.delimiter {
-                    tt::Delimiter::None => 0,
-                    _ => 2,
-                };
-
-                while let Some(tt) = iter.next() {
-                    count += Self::count_tt_tokens(&tt, Some(&mut iter));
-                }
-                count
-            }
-
-            tt::TokenTree::Leaf(leaf) => {
-                iter.map(|iter| {
-                    Self::skip_sibling_leaf(leaf, iter);
-                });
+fn convert_ident(ident: &tt::Ident) -> TtToken {
+    let kind = SyntaxKind::from_keyword(ident.text.as_str()).unwrap_or(IDENT);
+    TtToken { kind, is_joint_to_next: false, text: ident.text.clone(), n_tokens: 1 }
+}
 
-                1
-            }
-        }
+fn convert_punct(p: &tt::Punct, parent: &tt::Subtree, next: usize) -> TtToken {
+    let iter = parent.token_trees[next..].iter();
+    let mut peek = TokenPeek::new(iter);
+
+    if let Some((kind, is_joint_to_next, text, size)) = convert_multi_char_punct(p, &mut peek) {
+        TtToken { kind, is_joint_to_next, text: text.into(), n_tokens: size }
+    } else {
+        let kind = match p.char {
+            // lexer may produce combpund tokens for these ones
+            '.' => DOT,
+            ':' => COLON,
+            '=' => EQ,
+            '!' => EXCL,
+            '-' => MINUS,
+            c => SyntaxKind::from_char(c).unwrap(),
+        };
+        let text = {
+            let mut buf = [0u8; 4];
+            let s: &str = p.char.encode_utf8(&mut buf);
+            SmolStr::new(s)
+        };
+        TtToken { kind, is_joint_to_next: p.spacing == tt::Spacing::Joint, text, n_tokens: 1 }
     }
 }
diff --git a/crates/ra_mbe/src/syntax_bridge.rs b/crates/ra_mbe/src/syntax_bridge.rs
index b0fb91a63..19c17bd55 100644
--- a/crates/ra_mbe/src/syntax_bridge.rs
+++ b/crates/ra_mbe/src/syntax_bridge.rs
@@ -105,16 +105,16 @@ fn convert_tt(
     Some(res)
 }
 
-struct TtTreeSink<'a> {
+struct TtTreeSink<'a, Q: Querier> {
     buf: String,
-    src_querier: Querier<'a>,
+    src_querier: &'a Q,
     text_pos: TextUnit,
     token_pos: usize,
     inner: SyntaxTreeBuilder,
 }
 
-impl<'a> TtTreeSink<'a> {
-    fn new(src_querier: Querier<'a>) -> TtTreeSink {
+impl<'a, Q: Querier> TtTreeSink<'a, Q> {
+    fn new(src_querier: &'a Q) -> Self {
         TtTreeSink {
             buf: String::new(),
             src_querier,
@@ -125,10 +125,10 @@ impl<'a> TtTreeSink<'a> {
     }
 }
 
-impl<'a> TreeSink for TtTreeSink<'a> {
+impl<'a, Q: Querier> TreeSink for TtTreeSink<'a, Q> {
     fn token(&mut self, kind: SyntaxKind, n_tokens: u8) {
         for _ in 0..n_tokens {
-            self.buf += self.src_querier.token(self.token_pos).1;
+            self.buf += &self.src_querier.token(self.token_pos).1;
             self.token_pos += 1;
         }
         self.text_pos += TextUnit::of_str(&self.buf);