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
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

// micro-timer-macros
//
// Do not use this crate. It is only meant for internal use. Use `micro-timer`
// instead.
//
// Copyright 2020, Raphaël Gomès <rgomes@octobus.net>

/// `extern crate` Required even for 2018 edition
extern crate proc_macro;
use proc_macro2::TokenStream;
use quote::{quote, quote_spanned, ToTokens};
use syn::spanned::Spanned;

/// Logs the time elapsed for the body of the target function for each call.
///
/// ```compile_fail
/// use micro_timer::timed;
///
/// #[timed]  // Can only be used on functions
/// struct Thing;
/// ```
#[proc_macro_attribute]
pub fn timed(
    attrs: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    inner_timed(attrs.into(), item.into()).into()
}

/// This is the unit-testable version using `proc_macro2`
fn inner_timed(_attr_ts: TokenStream, fn_ts: TokenStream) -> TokenStream {
    let ast = syn::parse2(fn_ts.clone()).unwrap();
    let func = match parse_function(ast) {
        Ok(f) => f,
        Err(stream) => return stream,
    };

    let mut outer = func.clone();
    let original_func_name = func.sig.ident.to_string();
    let inner_block: TokenStream = func
        .block
        .stmts
        .iter()
        .map(|s| s.to_token_stream())
        .collect();
    let span = outer.sig.ident.span();

    // We use `__micro_timer_inner` as a name in case it pops up in error
    // messages to make it obvious that it's an issue caused by this crate,
    // it is not meant to prevent name collision.
    let block = quote_spanned! {
        span=>
        {
            let __micro_timer_instant = ::std::time::Instant::now();
            let __micro_timer_guard =
                ::micro_timer::scopeguard::guard_on_success(
                    __micro_timer_instant,
                    |timer| {
                        crate::log::trace!(
                            "Duration of `{}`: {:?}",
                            #original_func_name,
                            timer.elapsed()
                        );
                    }
                );

            #inner_block
        }
    };

    outer.block = syn::parse2(block).unwrap();

    (quote! {#outer})
}

fn parse_function(item: syn::Item) -> Result<syn::ItemFn, TokenStream> {
    match item {
        syn::Item::Fn(func) => Ok(func),
        i => Err(quote_spanned! {
            i.span()=>
            compile_error!("`#[timed]` can only be used on functions");
            #i
        }),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    #[test]
    fn test_output() {
        let input = syn::parse_quote! {
            fn my_super_function(_value: usize) -> usize {
                let timer = 10;
                timer
            }
        };

        let expected: TokenStream = syn::parse_quote! {
            fn my_super_function(_value: usize) -> usize {
                let __micro_timer_instant = ::std::time::Instant::now();
                ::micro_timer::scopeguard::guard_on_success(
                    __micro_timer_instant,
                    |timer| {
                        crate::log::trace!(
                            "Duration of `{}`: {:?}",
                            "my_super_function",
                            timer.elapsed()
                        );
                    }
                );
                let timer = 10;
                timer
            }
        };
        let output = inner_timed(TokenStream::new(), input);
        assert_eq!(output.to_string(), expected.to_string());
    }
}