6 files changed, 357 insertions, 70 deletions

diff --git a/crates/mbe/src/benchmark.rs b/crates/mbe/src/benchmark.rs
new file mode 100644
index 000000000..0d0acd589
--- /dev/null
+++ b/crates/mbe/src/benchmark.rs
@@ -0,0 +1,211 @@
+//! This module add real world mbe example for benchmark tests
+
+use rustc_hash::FxHashMap;
+use syntax::{
+    ast::{self, NameOwner},
+    AstNode, SmolStr,
+};
+use test_utils::{bench, bench_fixture, skip_slow_tests};
+
+use crate::{
+    ast_to_token_tree,
+    parser::{Op, RepeatKind, Separator},
+    MacroRules,
+};
+
+#[test]
+fn benchmark_parse_macro_rules() {
+    if skip_slow_tests() {
+        return;
+    }
+    let rules = macro_rules_fixtures_tt();
+    let hash: usize = {
+        let _pt = bench("mbe parse macro rules");
+        rules.values().map(|it| MacroRules::parse(it).unwrap().rules.len()).sum()
+    };
+    assert_eq!(hash, 1144);
+}
+
+#[test]
+fn benchmark_expand_macro_rules() {
+    if skip_slow_tests() {
+        return;
+    }
+    let rules = macro_rules_fixtures();
+    let invocations = invocation_fixtures(&rules);
+
+    let hash: usize = {
+        let _pt = bench("mbe expand macro rules");
+        invocations
+            .into_iter()
+            .map(|(id, tt)| {
+                let res = rules[&id].expand(&tt);
+                if res.err.is_some() {
+                    // FIXME:
+                    // Currently `invocation_fixtures` will generate some correct invocations but
+                    // cannot be expanded by mbe. We ignore errors here.
+                    // See: https://github.com/rust-analyzer/rust-analyzer/issues/4777
+                    eprintln!("err from {} {:?}", id, res.err);
+                }
+                res.value.token_trees.len()
+            })
+            .sum()
+    };
+    assert_eq!(hash, 66995);
+}
+
+fn macro_rules_fixtures() -> FxHashMap<String, MacroRules> {
+    macro_rules_fixtures_tt()
+        .into_iter()
+        .map(|(id, tt)| (id, MacroRules::parse(&tt).unwrap()))
+        .collect()
+}
+
+fn macro_rules_fixtures_tt() -> FxHashMap<String, tt::Subtree> {
+    let fixture = bench_fixture::numerous_macro_rules();
+    let source_file = ast::SourceFile::parse(&fixture).ok().unwrap();
+
+    source_file
+        .syntax()
+        .descendants()
+        .filter_map(ast::MacroRules::cast)
+        .map(|rule| {
+            let id = rule.name().unwrap().to_string();
+            let (def_tt, _) = ast_to_token_tree(&rule.token_tree().unwrap()).unwrap();
+            (id, def_tt)
+        })
+        .collect()
+}
+
+// Generate random invocation fixtures from rules
+fn invocation_fixtures(rules: &FxHashMap<String, MacroRules>) -> Vec<(String, tt::Subtree)> {
+    let mut seed = 123456789;
+    let mut res = Vec::new();
+
+    for (name, it) in rules {
+        for rule in &it.rules {
+            // Generate twice
+            for _ in 0..2 {
+                let mut subtree = tt::Subtree::default();
+                for op in rule.lhs.iter() {
+                    collect_from_op(op, &mut subtree, &mut seed);
+                }
+                res.push((name.clone(), subtree));
+            }
+        }
+    }
+    return res;
+
+    fn collect_from_op(op: &Op, parent: &mut tt::Subtree, seed: &mut usize) {
+        return match op {
+            Op::Var { kind, .. } => match kind.as_ref().map(|it| it.as_str()) {
+                Some("ident") => parent.token_trees.push(make_ident("foo")),
+                Some("ty") => parent.token_trees.push(make_ident("Foo")),
+                Some("tt") => parent.token_trees.push(make_ident("foo")),
+                Some("vis") => parent.token_trees.push(make_ident("pub")),
+                Some("pat") => parent.token_trees.push(make_ident("foo")),
+                Some("path") => parent.token_trees.push(make_ident("foo")),
+                Some("literal") => parent.token_trees.push(make_literal("1")),
+                Some("expr") => parent.token_trees.push(make_ident("foo").into()),
+                Some("lifetime") => {
+                    parent.token_trees.push(make_punct('\''));
+                    parent.token_trees.push(make_ident("a"));
+                }
+                Some("block") => {
+                    parent.token_trees.push(make_subtree(tt::DelimiterKind::Brace, None))
+                }
+                Some("item") => {
+                    parent.token_trees.push(make_ident("fn"));
+                    parent.token_trees.push(make_ident("foo"));
+                    parent.token_trees.push(make_subtree(tt::DelimiterKind::Parenthesis, None));
+                    parent.token_trees.push(make_subtree(tt::DelimiterKind::Brace, None));
+                }
+                Some("meta") => {
+                    parent.token_trees.push(make_ident("foo"));
+                    parent.token_trees.push(make_subtree(tt::DelimiterKind::Parenthesis, None));
+                }
+
+                None => (),
+                Some(kind) => panic!("Unhandled kind {}", kind),
+            },
+            Op::Leaf(leaf) => parent.token_trees.push(leaf.clone().into()),
+            Op::Repeat { tokens, kind, separator } => {
+                let max = 10;
+                let cnt = match kind {
+                    RepeatKind::ZeroOrMore => rand(seed) % max,
+                    RepeatKind::OneOrMore => 1 + rand(seed) % max,
+                    RepeatKind::ZeroOrOne => rand(seed) % 2,
+                };
+                for i in 0..cnt {
+                    for it in tokens.iter() {
+                        collect_from_op(it, parent, seed);
+                    }
+                    if i + 1 != cnt {
+                        if let Some(sep) = separator {
+                            match sep {
+                                Separator::Literal(it) => parent
+                                    .token_trees
+                                    .push(tt::Leaf::Literal(it.clone().into()).into()),
+                                Separator::Ident(it) => parent
+                                    .token_trees
+                                    .push(tt::Leaf::Ident(it.clone().into()).into()),
+                                Separator::Puncts(puncts) => {
+                                    for it in puncts {
+                                        parent
+                                            .token_trees
+                                            .push(tt::Leaf::Punct(it.clone().into()).into())
+                                    }
+                                }
+                            };
+                        }
+                    }
+                }
+            }
+            Op::Subtree { tokens, delimiter } => {
+                let mut subtree =
+                    tt::Subtree { delimiter: delimiter.clone(), token_trees: Vec::new() };
+                tokens.iter().for_each(|it| {
+                    collect_from_op(it, &mut subtree, seed);
+                });
+                parent.token_trees.push(subtree.into());
+            }
+        };
+
+        // Simple linear congruential generator for determistic result
+        fn rand(seed: &mut usize) -> usize {
+            let a = 1664525;
+            let c = 1013904223;
+            *seed = usize::wrapping_add(usize::wrapping_mul(*seed, a), c);
+            return *seed;
+        };
+        fn make_ident(ident: &str) -> tt::TokenTree {
+            tt::Leaf::Ident(tt::Ident { id: tt::TokenId::unspecified(), text: SmolStr::new(ident) })
+                .into()
+        }
+        fn make_punct(char: char) -> tt::TokenTree {
+            tt::Leaf::Punct(tt::Punct {
+                id: tt::TokenId::unspecified(),
+                char,
+                spacing: tt::Spacing::Alone,
+            })
+            .into()
+        }
+        fn make_literal(lit: &str) -> tt::TokenTree {
+            tt::Leaf::Literal(tt::Literal {
+                id: tt::TokenId::unspecified(),
+                text: SmolStr::new(lit),
+            })
+            .into()
+        }
+        fn make_subtree(
+            kind: tt::DelimiterKind,
+            token_trees: Option<Vec<tt::TokenTree>>,
+        ) -> tt::TokenTree {
+            tt::Subtree {
+                delimiter: Some(tt::Delimiter { id: tt::TokenId::unspecified(), kind }),
+                token_trees: token_trees.unwrap_or_default(),
+            }
+            .into()
+        }
+    }
+}
diff --git a/crates/mbe/src/expander/matcher.rs b/crates/mbe/src/expander/matcher.rs
index 800931cd1..e3bd4c09a 100644
--- a/crates/mbe/src/expander/matcher.rs
+++ b/crates/mbe/src/expander/matcher.rs
@@ -3,15 +3,13 @@
 use crate::{
     expander::{Binding, Bindings, Fragment},
     parser::{Op, RepeatKind, Separator},
-    subtree_source::SubtreeTokenSource,
     tt_iter::TtIter,
     ExpandError, MetaTemplate,
 };
 
 use super::ExpandResult;
-use parser::{FragmentKind::*, TreeSink};
-use syntax::{SmolStr, SyntaxKind};
-use tt::buffer::{Cursor, TokenBuffer};
+use parser::FragmentKind::*;
+use syntax::SmolStr;
 
 impl Bindings {
     fn push_optional(&mut self, name: &SmolStr) {
@@ -409,68 +407,6 @@ impl<'a> TtIter<'a> {
         .into())
     }
 
-    fn expect_fragment(
-        &mut self,
-        fragment_kind: parser::FragmentKind,
-    ) -> ExpandResult<Option<tt::TokenTree>> {
-        struct OffsetTokenSink<'a> {
-            cursor: Cursor<'a>,
-            error: bool,
-        }
-
-        impl<'a> TreeSink for OffsetTokenSink<'a> {
-            fn token(&mut self, kind: SyntaxKind, mut n_tokens: u8) {
-                if kind == SyntaxKind::LIFETIME_IDENT {
-                    n_tokens = 2;
-                }
-                for _ in 0..n_tokens {
-                    self.cursor = self.cursor.bump_subtree();
-                }
-            }
-            fn start_node(&mut self, _kind: SyntaxKind) {}
-            fn finish_node(&mut self) {}
-            fn error(&mut self, _error: parser::ParseError) {
-                self.error = true;
-            }
-        }
-
-        let buffer = TokenBuffer::from_tokens(&self.inner.as_slice());
-        let mut src = SubtreeTokenSource::new(&buffer);
-        let mut sink = OffsetTokenSink { cursor: buffer.begin(), error: false };
-
-        parser::parse_fragment(&mut src, &mut sink, fragment_kind);
-
-        let mut err = None;
-        if !sink.cursor.is_root() || sink.error {
-            err = Some(err!("expected {:?}", fragment_kind));
-        }
-
-        let mut curr = buffer.begin();
-        let mut res = vec![];
-
-        if sink.cursor.is_root() {
-            while curr != sink.cursor {
-                if let Some(token) = curr.token_tree() {
-                    res.push(token);
-                }
-                curr = curr.bump();
-            }
-        }
-        self.inner = self.inner.as_slice()[res.len()..].iter();
-        if res.len() == 0 && err.is_none() {
-            err = Some(err!("no tokens consumed"));
-        }
-        let res = match res.len() {
-            1 => Some(res[0].cloned()),
-            0 => None,
-            _ => Some(tt::TokenTree::Subtree(tt::Subtree {
-                delimiter: None,
-                token_trees: res.into_iter().map(|it| it.cloned()).collect(),
-            })),
-        };
-        ExpandResult { value: res, err }
-    }
-
     fn eat_vis(&mut self) -> Option<tt::TokenTree> {
         let mut fork = self.clone();
         match fork.expect_fragment(Visibility) {
diff --git a/crates/mbe/src/lib.rs b/crates/mbe/src/lib.rs
index d80bd7a33..4c298f85f 100644
--- a/crates/mbe/src/lib.rs
+++ b/crates/mbe/src/lib.rs
@@ -12,6 +12,9 @@ mod subtree_source;
 #[cfg(test)]
 mod tests;
 
+#[cfg(test)]
+mod benchmark;
+
 use std::fmt;
 
 use test_utils::mark;
@@ -62,8 +65,8 @@ impl fmt::Display for ExpandError {
 }
 
 pub use crate::syntax_bridge::{
-    ast_to_token_tree, parse_to_token_tree, syntax_node_to_token_tree, token_tree_to_syntax_node,
-    TokenMap,
+    ast_to_token_tree, parse_exprs_with_sep, parse_to_token_tree, syntax_node_to_token_tree,
+    token_tree_to_syntax_node, TokenMap,
 };
 
 /// This struct contains AST for a single `macro_rules` definition. What might
diff --git a/crates/mbe/src/syntax_bridge.rs b/crates/mbe/src/syntax_bridge.rs
index 0cdc175be..aacae1026 100644
--- a/crates/mbe/src/syntax_bridge.rs
+++ b/crates/mbe/src/syntax_bridge.rs
@@ -10,8 +10,8 @@ use syntax::{
 };
 use tt::buffer::{Cursor, TokenBuffer};
 
-use crate::subtree_source::SubtreeTokenSource;
 use crate::ExpandError;
+use crate::{subtree_source::SubtreeTokenSource, tt_iter::TtIter};
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub enum TokenTextRange {
@@ -112,6 +112,43 @@ pub fn parse_to_token_tree(text: &str) -> Option<(tt::Subtree, TokenMap)> {
     Some((subtree, conv.id_alloc.map))
 }
 
+/// Split token tree with seperate expr: $($e:expr)SEP*
+pub fn parse_exprs_with_sep(tt: &tt::Subtree, sep: char) -> Vec<tt::Subtree> {
+    if tt.token_trees.is_empty() {
+        return Vec::new();
+    }
+
+    let mut iter = TtIter::new(tt);
+    let mut res = Vec::new();
+
+    while iter.peek_n(0).is_some() {
+        let expanded = iter.expect_fragment(FragmentKind::Expr);
+        if expanded.err.is_some() {
+            break;
+        }
+
+        res.push(match expanded.value {
+            None => break,
+            Some(tt @ tt::TokenTree::Leaf(_)) => {
+                tt::Subtree { delimiter: None, token_trees: vec![tt.into()] }
+            }
+            Some(tt::TokenTree::Subtree(tt)) => tt,
+        });
+
+        let mut fork = iter.clone();
+        if fork.expect_char(sep).is_err() {
+            break;
+        }
+        iter = fork;
+    }
+
+    if iter.peek_n(0).is_some() {
+        res.push(tt::Subtree { delimiter: None, token_trees: iter.into_iter().cloned().collect() });
+    }
+
+    res
+}
+
 impl TokenMap {
     pub fn token_by_range(&self, relative_range: TextRange) -> Option<tt::TokenId> {
         let &(token_id, _) = self.entries.iter().find(|(_, range)| match range {
@@ -176,7 +213,7 @@ fn doc_comment_text(comment: &ast::Comment) -> SmolStr {
 
     // Quote the string
     // Note that `tt::Literal` expect an escaped string
-    let text = format!("{:?}", text.escape_default().to_string());
+    let text = format!("{:?}", text.escape_debug().to_string());
     text.into()
 }
 
diff --git a/crates/mbe/src/tests.rs b/crates/mbe/src/tests.rs
index 1c467facd..f1eadcd1e 100644
--- a/crates/mbe/src/tests.rs
+++ b/crates/mbe/src/tests.rs
@@ -970,6 +970,29 @@ fn test_meta_doc_comments() {
 }
 
 #[test]
+fn test_meta_doc_comments_non_latin() {
+    parse_macro(
+        r#"
+        macro_rules! foo {
+            ($(#[$ i:meta])+) => (
+                $(#[$ i])+
+                fn bar() {}
+            )
+        }
+"#,
+    ).
+    assert_expand_items(
+        r#"foo! {
+            /// 錦瑟無端五十弦，一弦一柱思華年。
+            /**
+                莊生曉夢迷蝴蝶，望帝春心託杜鵑。
+            */
+        }"#,
+        "# [doc = \" 錦瑟無端五十弦，一弦一柱思華年。\"] # [doc = \"\\\\n                莊生曉夢迷蝴蝶，望帝春心託杜鵑。\\\\n            \"] fn bar () {}",
+    );
+}
+
+#[test]
 fn test_tt_block() {
     parse_macro(
         r#"
diff --git a/crates/mbe/src/tt_iter.rs b/crates/mbe/src/tt_iter.rs
index 46c420718..a362d31fc 100644
--- a/crates/mbe/src/tt_iter.rs
+++ b/crates/mbe/src/tt_iter.rs
@@ -1,5 +1,20 @@
 //! FIXME: write short doc here
 
+use crate::{subtree_source::SubtreeTokenSource, ExpandError, ExpandResult};
+
+use parser::TreeSink;
+use syntax::SyntaxKind;
+use tt::buffer::{Cursor, TokenBuffer};
+
+macro_rules! err {
+    () => {
+        ExpandError::BindingError(format!(""))
+    };
+    ($($tt:tt)*) => {
+        ExpandError::BindingError(format!($($tt)*))
+    };
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct TtIter<'a> {
     pub(crate) inner: std::slice::Iter<'a, tt::TokenTree>,
@@ -56,6 +71,68 @@ impl<'a> TtIter<'a> {
         }
     }
 
+    pub(crate) fn expect_fragment(
+        &mut self,
+        fragment_kind: parser::FragmentKind,
+    ) -> ExpandResult<Option<tt::TokenTree>> {
+        struct OffsetTokenSink<'a> {
+            cursor: Cursor<'a>,
+            error: bool,
+        }
+
+        impl<'a> TreeSink for OffsetTokenSink<'a> {
+            fn token(&mut self, kind: SyntaxKind, mut n_tokens: u8) {
+                if kind == SyntaxKind::LIFETIME_IDENT {
+                    n_tokens = 2;
+                }
+                for _ in 0..n_tokens {
+                    self.cursor = self.cursor.bump_subtree();
+                }
+            }
+            fn start_node(&mut self, _kind: SyntaxKind) {}
+            fn finish_node(&mut self) {}
+            fn error(&mut self, _error: parser::ParseError) {
+                self.error = true;
+            }
+        }
+
+        let buffer = TokenBuffer::from_tokens(&self.inner.as_slice());
+        let mut src = SubtreeTokenSource::new(&buffer);
+        let mut sink = OffsetTokenSink { cursor: buffer.begin(), error: false };
+
+        parser::parse_fragment(&mut src, &mut sink, fragment_kind);
+
+        let mut err = None;
+        if !sink.cursor.is_root() || sink.error {
+            err = Some(err!("expected {:?}", fragment_kind));
+        }
+
+        let mut curr = buffer.begin();
+        let mut res = vec![];
+
+        if sink.cursor.is_root() {
+            while curr != sink.cursor {
+                if let Some(token) = curr.token_tree() {
+                    res.push(token);
+                }
+                curr = curr.bump();
+            }
+        }
+        self.inner = self.inner.as_slice()[res.len()..].iter();
+        if res.len() == 0 && err.is_none() {
+            err = Some(err!("no tokens consumed"));
+        }
+        let res = match res.len() {
+            1 => Some(res[0].cloned()),
+            0 => None,
+            _ => Some(tt::TokenTree::Subtree(tt::Subtree {
+                delimiter: None,
+                token_trees: res.into_iter().map(|it| it.cloned()).collect(),
+            })),
+        };
+        ExpandResult { value: res, err }
+    }
+
     pub(crate) fn peek_n(&self, n: usize) -> Option<&tt::TokenTree> {
         self.inner.as_slice().get(n)
     }