1use std::collections::BTreeMap;
2
3use proc_macro::TokenStream;
4use quote::{__private::Span, quote};
5use syn::{
6 parse::Parse, parse_macro_input, punctuated::Punctuated, Error, Ident, Lit, Meta,
7 MetaNameValue, Token,
8};
9use syn::{parse::ParseStream, ItemFn};
10use syn::{parse::Result, NestedMeta};
11
12#[derive(Debug)]
13struct Args {
14 attributes: BTreeMap<String, String>,
15}
16
17impl Parse for Args {
18 fn parse(input: ParseStream) -> Result<Self> {
19 try_parse(input)
20 }
21}
22
23fn try_parse(input: ParseStream) -> Result<Args> {
24 let args = Punctuated::<NestedMeta, Token![,]>::parse_terminated(input)?;
25
26 let attributes = args
27 .iter()
28 .map(|arg| match arg {
29 NestedMeta::Meta(Meta::NameValue(MetaNameValue {
30 lit: Lit::Str(lit_str),
31 path,
32 ..
33 })) => (
34 path.segments.first().unwrap().ident.to_string(),
35 lit_str.value(),
36 ),
37 _ => panic!("illegal arguments, example: output=\"eprintln\""),
38 })
39 .collect::<BTreeMap<_, _>>();
40 if attributes.get("output").map(|x| x.trim().is_empty()) == Some(true) {
41 return Err(error("output must be specified"));
42 }
43 Ok(Args { attributes })
44}
45
46fn error(msg: &str) -> Error {
47 Error::new(Span::call_site(), msg)
48}
49
50#[proc_macro_attribute]
51pub fn elapsed(args: TokenStream, func: TokenStream) -> TokenStream {
52 let args: Args = parse_macro_input!(args as Args);
53
54 let output_target = match args.attributes.get("output") {
55 Some(s) => s.to_string(),
56 _ => "println".to_string(),
57 };
58
59 let output_target = Ident::new(output_target.as_str(), Span::call_site());
60
61 let func = parse_macro_input!(func as ItemFn);
62 let vis = &func.vis;
63 let block = &func.block;
64
65 let signature = func.sig;
66 let asyncness = &signature.asyncness;
67 let constness = &signature.constness;
68 let unsafety = &signature.unsafety;
69 let abi = &signature.abi;
70 let ident = &signature.ident;
71 let name_str = ident.to_string();
72 let generics = &signature.generics;
73 let inputs = &signature.inputs;
74 let output = &signature.output;
75 let where_clause = &signature.generics.where_clause;
76
77 let new_fn = quote! {
78
79 #vis #constness #asyncness #unsafety #abi fn #ident #generics(#inputs) #output #where_clause {
80 use std::time;
81
82 let start = time::Instant::now();
83 let fn_return_value = #block;
84 #output_target!("fn {} costs {}", #name_str, human_time::human_time(start.elapsed()));
85 fn_return_value
86 }
87 };
88
89 new_fn.into()
90}