2 files changed, 80 insertions, 32 deletions

diff --git a/crates/ra_mbe/src/syntax_bridge.rs b/crates/ra_mbe/src/syntax_bridge.rs
index e0f228ce9..3521b382a 100644
--- a/crates/ra_mbe/src/syntax_bridge.rs
+++ b/crates/ra_mbe/src/syntax_bridge.rs
@@ -148,30 +148,21 @@ fn convert_tt(
         match child {
             SyntaxElement::Token(token) => {
                 if token.kind().is_punct() {
-                    let mut prev = None;
-                    for char in token.text().chars() {
-                        if let Some(char) = prev {
-                            token_trees.push(
-                                tt::Leaf::from(tt::Punct { char, spacing: tt::Spacing::Joint })
-                                    .into(),
-                            );
-                        }
-                        prev = Some(char)
-                    }
-                    if let Some(char) = prev {
-                        let spacing = match child_iter.peek() {
-                            Some(SyntaxElement::Token(token)) => {
-                                if token.kind().is_punct() {
-                                    tt::Spacing::Joint
-                                } else {
-                                    tt::Spacing::Alone
-                                }
+                    assert!(token.text().len() == 1, "Input ast::token punct must be single char.");
+                    let char = token.text().chars().next().unwrap();
+
+                    let spacing = match child_iter.peek() {
+                        Some(SyntaxElement::Token(token)) => {
+                            if token.kind().is_punct() {
+                                tt::Spacing::Joint
+                            } else {
+                                tt::Spacing::Alone
                             }
-                            _ => tt::Spacing::Alone,
-                        };
+                        }
+                        _ => tt::Spacing::Alone,
+                    };
 
-                        token_trees.push(tt::Leaf::from(tt::Punct { char, spacing }).into());
-                    }
+                    token_trees.push(tt::Leaf::from(tt::Punct { char, spacing }).into());
                 } else {
                     let child: tt::TokenTree = if token.kind() == SyntaxKind::TRUE_KW
                         || token.kind() == SyntaxKind::FALSE_KW
@@ -224,6 +215,15 @@ impl<'a, Q: Querier> TtTreeSink<'a, Q> {
     }
 }
 
+fn is_delimiter(kind: SyntaxKind) -> bool {
+    use SyntaxKind::*;
+
+    match kind {
+        L_PAREN | L_BRACK | L_CURLY | R_PAREN | R_BRACK | R_CURLY => true,
+        _ => false,
+    }
+}
+
 impl<'a, Q: Querier> TreeSink for TtTreeSink<'a, Q> {
     fn token(&mut self, kind: SyntaxKind, n_tokens: u8) {
         if kind == L_DOLLAR || kind == R_DOLLAR {
@@ -240,14 +240,18 @@ impl<'a, Q: Querier> TreeSink for TtTreeSink<'a, Q> {
         self.buf.clear();
         self.inner.token(kind, text);
 
-        // // Add a white space to token
-        // let (last_kind, _, last_joint_to_next ) = self.src_querier.token(self.token_pos-n_tokens as usize);
-        // if !last_joint_to_next && last_kind.is_punct() {
-        //     let (cur_kind, _, _ ) = self.src_querier.token(self.token_pos);
-        //     if cur_kind.is_punct() {
-        //         self.inner.token(WHITESPACE, " ".into());
-        //     }
-        // }
+        // Add a white space between tokens, only if both are not delimiters
+        if !is_delimiter(kind) {
+            let (last_kind, _, last_joint_to_next) = self.src_querier.token(self.token_pos - 1);
+            if !last_joint_to_next && last_kind.is_punct() {
+                let (cur_kind, _, _) = self.src_querier.token(self.token_pos);
+                if !is_delimiter(cur_kind) {
+                    if cur_kind.is_punct() {
+                        self.inner.token(WHITESPACE, " ".into());
+                    }
+                }
+            }
+        }
     }
 
     fn start_node(&mut self, kind: SyntaxKind) {
diff --git a/crates/ra_mbe/src/tt_cursor.rs b/crates/ra_mbe/src/tt_cursor.rs
index eef642a9c..53bc305be 100644
--- a/crates/ra_mbe/src/tt_cursor.rs
+++ b/crates/ra_mbe/src/tt_cursor.rs
@@ -1,6 +1,5 @@
 use crate::ParseError;
 use crate::subtree_parser::Parser;
-use crate::subtree_source::TokenPeek;
 use smallvec::{SmallVec, smallvec};
 
 #[derive(Debug, Clone)]
@@ -153,7 +152,7 @@ impl<'a> TtCursor<'a> {
     pub(crate) fn eat_vis(&mut self) -> Option<tt::TokenTree> {
         let parser = Parser::new(&mut self.pos, self.subtree);
         parser.parse_vis()
-    }
+        }
 
     pub(crate) fn expect_char(&mut self, char: char) -> Result<(), ParseError> {
         if self.at_char(char) {
@@ -262,3 +261,48 @@ impl<'a> TtCursor<'a> {
         self.pos = memento.pos;
     }
 }
+
+pub(crate) struct TokenPeek<'a, I>
+where
+    I: Iterator<Item = &'a tt::TokenTree>,
+{
+    iter: itertools::MultiPeek<I>,
+}
+
+// helper function
+fn to_punct(tt: &tt::TokenTree) -> Option<&tt::Punct> {
+    if let tt::TokenTree::Leaf(tt::Leaf::Punct(pp)) = tt {
+        return Some(pp);
+    }
+    None
+}
+
+impl<'a, I> TokenPeek<'a, I>
+where
+    I: Iterator<Item = &'a tt::TokenTree>,
+{
+    pub fn new(iter: I) -> Self {
+        TokenPeek { iter: itertools::multipeek(iter) }
+    }
+
+    pub fn current_punct2(&mut self, p: &tt::Punct) -> Option<((char, char), bool)> {
+        if p.spacing != tt::Spacing::Joint {
+            return None;
+        }
+
+        self.iter.reset_peek();
+        let p1 = to_punct(self.iter.peek()?)?;
+        Some(((p.char, p1.char), p1.spacing == tt::Spacing::Joint))
+    }
+
+    pub fn current_punct3(&mut self, p: &tt::Punct) -> Option<((char, char, char), bool)> {
+        self.current_punct2(p).and_then(|((p0, p1), last_joint)| {
+            if !last_joint {
+                None
+            } else {
+                let p2 = to_punct(*self.iter.peek()?)?;
+                Some(((p0, p1, p2.char), p2.spacing == tt::Spacing::Joint))
+            }
+        })
+    }
+}