elapsed_time/
lib.rs

1use proc_macro::TokenStream;
2use quote::{quote, ToTokens};
3use syn::{parse_macro_input, punctuated::Punctuated, Block, ItemFn, Meta};
4
5type AttributeArgs = Punctuated<Meta, syn::token::Comma>;
6
7#[proc_macro_attribute]
8pub fn elapsed(_args: TokenStream, function_def: TokenStream) -> TokenStream {
9    let mut item = syn::parse(function_def).unwrap();
10    let fn_item = match &mut item {
11        syn::Item::Fn(fn_item) => fn_item,
12        _ => panic!("expected fn")
13    };
14    let ItemFn { attrs, vis, sig, block } = fn_item;
15    let function_body = block.clone();
16    let fn_name = sig.ident.clone();
17    let log_ns = format!("{} took {{}}ns", fn_name);
18    let log_us = format!("{} took {{}}µs", fn_name);
19    let log_ms = format!("{} took {{}}ms", fn_name);
20    #[cfg(feature = "tracing")]
21    let log_ns_stmt = quote! { tracing::debug!(#log_ns, elapsed); };
22    #[cfg(all(not(feature = "tracing"), feature = "log"))]
23    let log_ns_stmt = quote! { log::debug!(#log_ns, elapsed); };
24    #[cfg(all(not(feature = "tracing"), not(feature = "log")))]
25    let log_ns_stmt = quote! { println!(#log_ns, elapsed); };
26
27    #[cfg(feature = "tracing")]
28    let log_us_stmt = quote! { tracing::debug!(#log_us, elapsed as f64 / 1000.0); };
29    #[cfg(all(not(feature = "tracing"), feature = "log"))]
30    let log_us_stmt = quote! { log::debug!(#log_us, elapsed as f64 / 1000.0); };
31    #[cfg(all(not(feature = "tracing"), not(feature = "log")))]
32    let log_us_stmt = quote! { println!(#log_us, elapsed as f64 / 1000.0); };
33
34    #[cfg(feature = "tracing")]
35    let log_ms_stmt = quote! { tracing::debug!(#log_ms, elapsed as f64 / 1000.0 / 1000.0); };
36    #[cfg(all(not(feature = "tracing"), feature = "log"))]
37    let log_ms_stmt = quote! { log::debug!(#log_ms, elapsed as f64 / 1000.0 / 1000.0); };
38    #[cfg(all(not(feature = "tracing"), not(feature = "log")))]
39    let log_ms_stmt = quote! { println!(#log_ms, elapsed as f64 / 1000.0 / 1000.0); };
40    let new_function_def = quote! {
41        #(#attrs)* #vis #sig {
42            let start_for_elapsed_macro = std::time::Instant::now();
43            let mut wrapped_func = || #function_body;
44            let res = wrapped_func();
45            let elapsed = start_for_elapsed_macro.elapsed().as_nanos();
46            if elapsed < 1000 {
47                #log_ns_stmt
48            } else if elapsed < 1000 * 1000 {
49                #log_us_stmt
50            } else {
51                #log_ms_stmt
52            }
53            res
54        }
55    };
56    TokenStream::from(new_function_def)
57}
58
59#[proc_macro_attribute]
60pub fn elapsed_block(args: TokenStream, block_def: TokenStream) -> TokenStream {
61    let mut block_name = "unnamed block".to_string();
62    let attrs = parse_macro_input!(args with AttributeArgs::parse_terminated);
63    if let Some(first_attr) = attrs.first() {
64        block_name = first_attr.to_token_stream().to_string();
65    }
66    let item = syn::parse::<Block>(block_def).unwrap();
67    let log_ns = format!("{} took {{}}ns", block_name);
68    let log_us = format!("{} took {{}}µs", block_name);
69    let log_ms = format!("{} took {{}}ms", block_name);
70    #[cfg(feature = "tracing")]
71    let log_ns_stmt = quote! { tracing::debug!(#log_ns, elapsed); };
72    #[cfg(all(not(feature = "tracing"), feature = "log"))]
73    let log_ns_stmt = quote! { log::debug!(#log_ns, elapsed); };
74    #[cfg(all(not(feature = "tracing"), not(feature = "log")))]
75    let log_ns_stmt = quote! { println!(#log_ns, elapsed); };
76
77    #[cfg(feature = "tracing")]
78    let log_us_stmt = quote! { tracing::debug!(#log_us, elapsed as f64 / 1000.0); };
79    #[cfg(all(not(feature = "tracing"), feature = "log"))]
80    let log_us_stmt = quote! { log::debug!(#log_us, elapsed as f64 / 1000.0); };
81    #[cfg(all(not(feature = "tracing"), not(feature = "log")))]
82    let log_us_stmt = quote! { println!(#log_us, elapsed as f64 / 1000.0); };
83
84    #[cfg(feature = "tracing")]
85    let log_ms_stmt = quote! { tracing::debug!(#log_ms, elapsed as f64 / 1000.0 / 1000.0); };
86    #[cfg(all(not(feature = "tracing"), feature = "log"))]
87    let log_ms_stmt = quote! { log::debug!(#log_ms, elapsed as f64 / 1000.0 / 1000.0); };
88    #[cfg(all(not(feature = "tracing"), not(feature = "log")))]
89    let log_ms_stmt = quote! { println!(#log_ms, elapsed as f64 / 1000.0 / 1000.0); };
90    let new_block_def = quote! {
91        {
92            let start_for_elapsed_macro = std::time::Instant::now();
93            #item
94            let elapsed = start_for_elapsed_macro.elapsed().as_nanos();
95            if elapsed < 1000 {
96                #log_ns_stmt
97            } else if elapsed < 1000 * 1000 {
98                #log_us_stmt
99            } else {
100                #log_ms_stmt
101            }
102        }
103    };
104    TokenStream::from(new_block_def)
105}
elapsed_time/lib.rs

elapsed_time/
lib.rs
