1extern crate proc_macro;
2use proc_macro::TokenStream;
3use proc_macro2::{Ident, TokenTree};
4use quote::quote_spanned;
5
6#[proc_macro_attribute]
7pub fn criterion(attr: TokenStream, item: TokenStream) -> TokenStream {
8 let attr = proc_macro2::TokenStream::from(attr);
9 let item = proc_macro2::TokenStream::from(item);
10
11 let span = proc_macro2::Span::call_site();
12
13 let init = if stream_length(attr.clone()) != 0 {
14 attr
15 }
16 else {
17 quote_spanned!(span=> criterion::Criterion::default())
18 };
19
20 let function_name = find_name(item.clone());
21 let wrapped_name = Ident::new(&format!("criterion_wrapped_{}", function_name.to_string()), span);
22
23 let output = quote_spanned!(span=>
24 #[test_case]
25 pub fn #wrapped_name() {
26 #item
27
28 let mut c = #init.configure_from_args();
29 #function_name(&mut c);
30 }
31 );
32
33 output.into()
34}
35
36fn stream_length(stream: proc_macro2::TokenStream) -> usize {
37 stream.into_iter().count()
38}
39
40fn find_name(stream: proc_macro2::TokenStream) -> Ident {
41 let mut iter = stream.into_iter();
42 while let Some(tok) = iter.next() {
43 if let TokenTree::Ident(ident) = tok {
44 if ident == "fn" {
45 break;
46 }
47 }
48 }
49
50 if let Some(TokenTree::Ident(name)) = iter.next() {
51 name
52 }
53 else {
54 panic!("Unable to find function name")
55 }
56}