crates/mbe/src/benchmark.rs

//! This module add real world mbe example for benchmark tests

use intern::Symbol;
use rustc_hash::FxHashMap;
use span::{Edition, Span};
use stdx::itertools::Itertools;
use syntax::{
    AstNode,
    ast::{self, HasName},
};
use syntax_bridge::{
    DocCommentDesugarMode,
    dummy_test_span_utils::{DUMMY, DummyTestSpanMap},
    syntax_node_to_token_tree,
};
use test_utils::{bench, bench_fixture, skip_slow_tests};

use crate::{
    DeclarativeMacro,
    parser::{MetaVarKind, Op, RepeatKind, Separator},
};

#[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
            .into_iter()
            .sorted_by_key(|(id, _)| id.clone())
            .map(|(_, it)| {
                DeclarativeMacro::parse_macro_rules(&it, |_| span::Edition::CURRENT).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, |_| (), DUMMY, Edition::CURRENT);
                assert!(res.err.is_none());
                res.value.0.0.len()
            })
            .sum()
    };
    assert_eq!(hash, 450144);
}

fn macro_rules_fixtures() -> FxHashMap<String, DeclarativeMacro> {
    macro_rules_fixtures_tt()
        .into_iter()
        .sorted_by_key(|(id, _)| id.clone())
        .map(|(id, tt)| (id, DeclarativeMacro::parse_macro_rules(&tt, |_| span::Edition::CURRENT)))
        .collect()
}

fn macro_rules_fixtures_tt() -> FxHashMap<String, tt::TopSubtree<Span>> {
    let fixture = bench_fixture::numerous_macro_rules();
    let source_file = ast::SourceFile::parse(&fixture, span::Edition::CURRENT).ok().unwrap();

    source_file
        .syntax()
        .descendants()
        .filter_map(ast::MacroRules::cast)
        .map(|rule| {
            let id = rule.name().unwrap().to_string();
            let def_tt = syntax_node_to_token_tree(
                rule.token_tree().unwrap().syntax(),
                DummyTestSpanMap,
                DUMMY,
                DocCommentDesugarMode::Mbe,
            );
            (id, def_tt)
        })
        .collect()
}

/// Generate random invocation fixtures from rules
fn invocation_fixtures(
    rules: &FxHashMap<String, DeclarativeMacro>,
) -> Vec<(String, tt::TopSubtree<Span>)> {
    let mut seed = 123456789;
    let mut res = Vec::new();

    for (name, it) in rules.iter().sorted_by_key(|&(id, _)| id) {
        for rule in it.rules.iter() {
            // Generate twice
            for _ in 0..2 {
                // The input are generated by filling the `Op` randomly.
                // However, there are some cases generated are ambiguous for expanding, for example:
                // ```rust
                // macro_rules! m {
                //    ($($t:ident),* as $ty:ident) => {}
                // }
                // m!(as u32);  // error: local ambiguity: multiple parsing options: built-in NTs ident ('t') or 1 other option.
                // ```
                //
                // So we just skip any error cases and try again
                let mut try_cnt = 0;
                loop {
                    let mut builder = tt::TopSubtreeBuilder::new(tt::Delimiter {
                        open: DUMMY,
                        close: DUMMY,
                        kind: tt::DelimiterKind::Invisible,
                    });
                    for op in rule.lhs.iter() {
                        collect_from_op(op, &mut builder, &mut seed);
                    }
                    let subtree = builder.build();

                    if it.expand(&subtree, |_| (), DUMMY, Edition::CURRENT).err.is_none() {
                        res.push((name.clone(), subtree));
                        break;
                    }
                    try_cnt += 1;
                    if try_cnt > 100 {
                        panic!("invocation fixture {name} cannot be generated.\n");
                    }
                }
            }
        }
    }
    return res;

    fn collect_from_op(op: &Op, builder: &mut tt::TopSubtreeBuilder<Span>, seed: &mut usize) {
        return match op {
            Op::Var { kind, .. } => match kind.as_ref() {
                Some(MetaVarKind::Ident) => builder.push(make_ident("foo")),
                Some(MetaVarKind::Ty) => builder.push(make_ident("Foo")),
                Some(MetaVarKind::Tt) => builder.push(make_ident("foo")),
                Some(MetaVarKind::Vis) => builder.push(make_ident("pub")),
                Some(MetaVarKind::Pat) => builder.push(make_ident("foo")),
                Some(MetaVarKind::Path) => builder.push(make_ident("foo")),
                Some(MetaVarKind::Literal) => builder.push(make_literal("1")),
                Some(MetaVarKind::Expr(_)) => builder.push(make_ident("foo")),
                Some(MetaVarKind::Lifetime) => {
                    builder.push(make_punct('\''));
                    builder.push(make_ident("a"));
                }
                Some(MetaVarKind::Block) => make_subtree(tt::DelimiterKind::Brace, builder),
                Some(MetaVarKind::Item) => {
                    builder.push(make_ident("fn"));
                    builder.push(make_ident("foo"));
                    make_subtree(tt::DelimiterKind::Parenthesis, builder);
                    make_subtree(tt::DelimiterKind::Brace, builder);
                }
                Some(MetaVarKind::Meta) => {
                    builder.push(make_ident("foo"));
                    make_subtree(tt::DelimiterKind::Parenthesis, builder);
                }

                None => (),
                Some(kind) => panic!("Unhandled kind {kind:?}"),
            },
            Op::Literal(it) => builder.push(tt::Leaf::from(it.clone())),
            Op::Ident(it) => builder.push(tt::Leaf::from(it.clone())),
            Op::Punct(puncts) => {
                for punct in puncts.as_slice() {
                    builder.push(tt::Leaf::from(*punct));
                }
            }
            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, builder, seed);
                    }
                    if i + 1 != cnt {
                        if let Some(sep) = separator {
                            match &**sep {
                                Separator::Literal(it) => {
                                    builder.push(tt::Leaf::Literal(it.clone()))
                                }
                                Separator::Ident(it) => builder.push(tt::Leaf::Ident(it.clone())),
                                Separator::Puncts(puncts) => {
                                    for it in puncts {
                                        builder.push(tt::Leaf::Punct(*it))
                                    }
                                }
                            };
                        }
                    }
                }
            }
            Op::Subtree { tokens, delimiter } => {
                builder.open(delimiter.kind, delimiter.open);
                tokens.iter().for_each(|it| collect_from_op(it, builder, seed));
                builder.close(delimiter.close);
            }
            Op::Ignore { .. }
            | Op::Index { .. }
            | Op::Count { .. }
            | Op::Len { .. }
            | Op::Concat { .. } => {}
        };

        // Simple linear congruential generator for deterministic result
        fn rand(seed: &mut usize) -> usize {
            let a = 1664525;
            let c = 1013904223;
            *seed = usize::wrapping_add(usize::wrapping_mul(*seed, a), c);
            *seed
        }
        fn make_ident(ident: &str) -> tt::Leaf<Span> {
            tt::Leaf::Ident(tt::Ident {
                span: DUMMY,
                sym: Symbol::intern(ident),
                is_raw: tt::IdentIsRaw::No,
            })
        }
        fn make_punct(char: char) -> tt::Leaf<Span> {
            tt::Leaf::Punct(tt::Punct { span: DUMMY, char, spacing: tt::Spacing::Alone })
        }
        fn make_literal(lit: &str) -> tt::Leaf<Span> {
            tt::Leaf::Literal(tt::Literal {
                span: DUMMY,
                symbol: Symbol::intern(lit),
                kind: tt::LitKind::Str,
                suffix: None,
            })
        }
        fn make_subtree(kind: tt::DelimiterKind, builder: &mut tt::TopSubtreeBuilder<Span>) {
            builder.open(kind, DUMMY);
            builder.close(DUMMY);
        }
    }
}