1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

use proc_macro::TokenStream;
use quote::{quote, ToTokens};
use syn::{Block, ItemFn, parse_macro_input};
use syn::AttributeArgs;

#[proc_macro_attribute]
pub fn elapsed(_args: TokenStream, function_def: TokenStream) -> TokenStream {
    let mut item = syn::parse(function_def).unwrap();
    let fn_item = match &mut item {
        syn::Item::Fn(fn_item) => fn_item,
        _ => panic!("expected fn")
    };
    let ItemFn { attrs, vis, sig, block } = fn_item;
    let function_body = block.clone();
    let fn_name = sig.ident.clone();
    let log_ns = format!("{} took {{}}ns", fn_name);
    let log_us = format!("{} took {{}}µs", fn_name);
    let log_ms = format!("{} took {{}}ms", fn_name);
    let new_function_def = quote! {
        #(#attrs)* #vis #sig {
            let start_for_elapsed_macro = std::time::Instant::now();
            let mut wrapped_func = || #function_body;
            let res = wrapped_func();
            let elapsed = start_for_elapsed_macro.elapsed().as_nanos();
            if elapsed < 1000 {
                println!(#log_ns, elapsed);
            } else if elapsed < 1000 * 1000 {
                println!(#log_us, elapsed as f64 / 1000.0);
            } else {
                println!(#log_ms, elapsed as f64 / 1000.0 / 1000.0);
            }
            res
        }
    };
    TokenStream::from(new_function_def)
}

#[proc_macro_attribute]
pub fn elapsed_block(args: TokenStream, block_def: TokenStream) -> TokenStream {
    let mut block_name = "unnamed block".to_string();
    let attrs = parse_macro_input!(args as AttributeArgs);
    if attrs.len() > 0 {
        block_name = attrs.get(0).unwrap().to_token_stream().to_string();
    }
    let item = syn::parse::<Block>(block_def).unwrap();
    let log_ns = format!("{} took {{}}ns", block_name);
    let log_us = format!("{} took {{}}µs", block_name);
    let log_ms = format!("{} took {{}}ms", block_name);
    let new_block_def = quote! {
        {
            let start_for_elapsed_macro = std::time::Instant::now();
            #item
            let elapsed = start_for_elapsed_macro.elapsed().as_nanos();
            if elapsed < 1000 {
                println!(#log_ns, elapsed);
            } else if elapsed < 1000 * 1000 {
                println!(#log_us, elapsed as f64 / 1000.0);
            } else {
                println!(#log_ms, elapsed as f64 / 1000.0 / 1000.0);
            }
        }
    };
    TokenStream::from(new_block_def)
}