Merge pull request #3451 from edwin0cheng/fix-mbe-composited

Fix mbe composited token bug

4 files changed, 107 insertions, 4 deletions

diff --git a/crates/ra_mbe/src/mbe_expander/matcher.rs b/crates/ra_mbe/src/mbe_expander/matcher.rs
index 2bdea11e1..ffba03898 100644
--- a/crates/ra_mbe/src/mbe_expander/matcher.rs
+++ b/crates/ra_mbe/src/mbe_expander/matcher.rs
@@ -155,6 +155,60 @@ impl<'a> TtIter<'a> {
         ok
     }
 
+    pub(crate) fn expect_tt(&mut self) -> Result<tt::TokenTree, ()> {
+        let tt = self.next().ok_or_else(|| ())?.clone();
+        let punct = match tt {
+            tt::TokenTree::Leaf(tt::Leaf::Punct(punct)) if punct.spacing == tt::Spacing::Joint => {
+                punct
+            }
+            _ => return Ok(tt),
+        };
+
+        let (second, third) = match (self.peek_n(0), self.peek_n(1)) {
+            (
+                Some(tt::TokenTree::Leaf(tt::Leaf::Punct(p2))),
+                Some(tt::TokenTree::Leaf(tt::Leaf::Punct(p3))),
+            ) if p2.spacing == tt::Spacing::Joint => (p2.char, Some(p3.char)),
+            (Some(tt::TokenTree::Leaf(tt::Leaf::Punct(p2))), _) => (p2.char, None),
+            _ => return Ok(tt),
+        };
+
+        match (punct.char, second, third) {
+            ('.', '.', Some('.'))
+            | ('.', '.', Some('='))
+            | ('<', '<', Some('='))
+            | ('>', '>', Some('=')) => {
+                let tt2 = self.next().unwrap().clone();
+                let tt3 = self.next().unwrap().clone();
+                Ok(tt::Subtree { delimiter: None, token_trees: vec![tt, tt2, tt3] }.into())
+            }
+            ('-', '=', None)
+            | ('-', '>', None)
+            | (':', ':', None)
+            | ('!', '=', None)
+            | ('.', '.', None)
+            | ('*', '=', None)
+            | ('/', '=', None)
+            | ('&', '&', None)
+            | ('&', '=', None)
+            | ('%', '=', None)
+            | ('^', '=', None)
+            | ('+', '=', None)
+            | ('<', '<', None)
+            | ('<', '=', None)
+            | ('=', '=', None)
+            | ('=', '>', None)
+            | ('>', '=', None)
+            | ('>', '>', None)
+            | ('|', '=', None)
+            | ('|', '|', None) => {
+                let tt2 = self.next().unwrap().clone();
+                Ok(tt::Subtree { delimiter: None, token_trees: vec![tt.clone(), tt2] }.into())
+            }
+            _ => Ok(tt),
+        }
+    }
+
     pub(crate) fn expect_lifetime(&mut self) -> Result<&tt::Ident, ()> {
         let ident = self.expect_ident()?;
         // check if it start from "`"
@@ -302,7 +356,7 @@ fn match_meta_var(kind: &str, input: &mut TtIter) -> Result<Option<Fragment>, Ex
                     let ident = input.expect_ident().map_err(|()| err!("expected ident"))?.clone();
                     tt::Leaf::from(ident).into()
                 }
-                "tt" => input.next().ok_or_else(|| err!())?.clone(),
+                "tt" => input.expect_tt().map_err(|()| err!())?.clone(),
                 "lifetime" => {
                     let ident = input.expect_lifetime().map_err(|()| err!())?;
                     tt::Leaf::Ident(ident.clone()).into()
diff --git a/crates/ra_mbe/src/syntax_bridge.rs b/crates/ra_mbe/src/syntax_bridge.rs
index 2aaf0215f..d8ee74faa 100644
--- a/crates/ra_mbe/src/syntax_bridge.rs
+++ b/crates/ra_mbe/src/syntax_bridge.rs
@@ -388,11 +388,12 @@ impl<'a> TreeSink for TtTreeSink<'a> {
             return;
         }
 
+        let mut last = self.cursor;
         for _ in 0..n_tokens {
             if self.cursor.eof() {
                 break;
             }
-
+            last = self.cursor;
             let text: SmolStr = match self.cursor.token_tree() {
                 Some(tt::TokenTree::Leaf(leaf)) => {
                     // Mark the range if needed
@@ -441,11 +442,11 @@ impl<'a> TreeSink for TtTreeSink<'a> {
         self.inner.token(kind, text);
 
         // Add whitespace between adjoint puncts
-        let next = self.cursor.bump();
+        let next = last.bump();
         if let (
             Some(tt::TokenTree::Leaf(tt::Leaf::Punct(curr))),
             Some(tt::TokenTree::Leaf(tt::Leaf::Punct(_))),
-        ) = (self.cursor.token_tree(), next.token_tree())
+        ) = (last.token_tree(), next.token_tree())
         {
             if curr.spacing == tt::Spacing::Alone {
                 self.inner.token(WHITESPACE, " ".into());
diff --git a/crates/ra_mbe/src/tests.rs b/crates/ra_mbe/src/tests.rs
index 1dba82915..304867881 100644
--- a/crates/ra_mbe/src/tests.rs
+++ b/crates/ra_mbe/src/tests.rs
@@ -825,6 +825,50 @@ fn test_tt_group() {
     )
     .assert_expand_items(r#"foo! { fn foo() {} }"#, r#"fn foo () {}"#);
 }
+
+#[test]
+fn test_tt_composite() {
+    parse_macro(
+        r#"
+            macro_rules! foo {
+                 ($i:tt) => { 0 }
+            }
+    "#,
+    )
+    .assert_expand_items(r#"foo! { => }"#, r#"0"#);
+}
+
+#[test]
+fn test_tt_composite2() {
+    let node = parse_macro(
+        r#"
+            macro_rules! foo {
+                ($($tt:tt)*) => { abs!(=> $($tt)*) }
+            }
+    "#,
+    )
+    .expand_items(r#"foo!{#}"#);
+
+    let res = format!("{:#?}", &node);
+    assert_eq_text!(
+        res.trim(),
+        r###"MACRO_ITEMS@[0; 10)
+  MACRO_CALL@[0; 10)
+    PATH@[0; 3)
+      PATH_SEGMENT@[0; 3)
+        NAME_REF@[0; 3)
+          IDENT@[0; 3) "abs"
+    EXCL@[3; 4) "!"
+    TOKEN_TREE@[4; 10)
+      L_PAREN@[4; 5) "("
+      EQ@[5; 6) "="
+      R_ANGLE@[6; 7) ">"
+      WHITESPACE@[7; 8) " "
+      POUND@[8; 9) "#"
+      R_PAREN@[9; 10) ")""###
+    );
+}
+
 #[test]
 fn test_lifetime() {
     parse_macro(
diff --git a/crates/ra_mbe/src/tt_iter.rs b/crates/ra_mbe/src/tt_iter.rs
index 319f1ee65..100184e66 100644
--- a/crates/ra_mbe/src/tt_iter.rs
+++ b/crates/ra_mbe/src/tt_iter.rs
@@ -53,6 +53,10 @@ impl<'a> TtIter<'a> {
             _ => Err(()),
         }
     }
+
+    pub(crate) fn peek_n(&self, n: usize) -> Option<&tt::TokenTree> {
+        self.inner.as_slice().get(n)
+    }
 }
 
 impl<'a> Iterator for TtIter<'a> {