human_time_macros/
lib.rs

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}