1extern crate proc_macro;
2
3use std::str::FromStr;
4
5use darling::FromMeta;
6use proc_macro2::TokenStream;
7use quote::{quote, ToTokens};
8
9#[derive(Debug, FromMeta)]
10struct Args {
11 #[darling(default)]
12 name: Option<String>,
13
14 #[darling(default)]
15 count: Option<usize>,
16
17 #[darling(default)]
18 no_test: Option<()>
19}
20
21#[proc_macro_attribute]
22pub fn bench(attrs: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream {
23 let func = syn::parse_macro_input!(item as syn::ItemFn);
24 let func_name = &func.sig.ident;
25 let func_attrs = &func.attrs;
26
27
28 let args: Args = Args::from_list(&syn::parse_macro_input!(attrs as syn::AttributeArgs)).unwrap();
29 let name = args.name.map(|s| s.to_token_stream()).unwrap_or(func_name.to_string().to_token_stream());
30 let count = args.count.unwrap_or(1000).to_token_stream();
31 let test = if args.no_test.is_some() {
32 TokenStream::new()
33 } else {
34 TokenStream::from_str("#[test]").unwrap()
35 };
36
37 let bencher = if cfg!(feature = "track-allocator") {
38 quote! {
39 Bencher::new(#name, #count, 0)
40 }
41 } else {
42 quote! {
43 Bencher::new(#name, #count, 0, GLOBAL.counter(), GLOBAL.peak())
44 }
45 };
46
47 (quote! {
48 #test
49 #(#func_attrs)*
50 fn #func_name() {
51 #func
52
53 let mut bencher = #bencher;
54 #func_name(&mut bencher);
55 bencher.finish();
56 }
57 }).into()
58}