// micro-timer
//
// 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};
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();
    outer.sig.ident = func.sig.ident.to_owned();
    let original_func_name = func.sig.ident.to_string();

    // Figure out if the `inner` closure needs to be mut based on whether
    // the original signature contains a `mut self`, `&mut self`,
    // `mut thing: T` or a `thing: &mut T`.
    // There are probably more ways you can have a signature need mutability,
    // but this is good enough for now.
    let needs_mut = func.sig.inputs.iter().any(|arg| match arg {
        syn::FnArg::Receiver(rec) => rec.mutability.is_some(),
        syn::FnArg::Typed(typed) => match *typed.pat {
            syn::Pat::Ident(syn::PatIdent { mutability, .. }) => {
                // Is it `mut arg: T`?
                if mutability.is_some() {
                    return true;
                }
                // Is it `arg: &mut T`?
                match *typed.ty {
                    syn::Type::Reference(ref r) => r.mutability.is_some(),
                    _ => false,
                }
            }
            _ => false,
        },
    });

    let mutability = if needs_mut {
        quote! {mut}
    } else {
        quote! {}
    };

    let inner_block = &func.block;
    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 timer = ::std::time::Instant::now();
            let #mutability __micro_timer_inner = || #inner_block;
            let ret = __micro_timer_inner();

            crate::log::trace!(
                "Duration of `{}`: {:?}",
                #original_func_name,
                timer.elapsed()
            );

            ret
        }
    };

    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 timer = ::std::time::Instant::now();

                let __micro_timer_inner = | | {
                    let timer = 10;
                    timer
                };
                let ret = __micro_timer_inner();

                crate::log::trace!(
                    "Duration of `{}`: {:?}",
                    "my_super_function",
                    timer.elapsed()
                );

                ret
            }
        };
        let output = inner_timed(TokenStream::new(), input);
        assert_eq!(output.to_string(), expected.to_string());
    }

    #[test]
    fn test_mut_output() {
        let input = syn::parse_quote! {
            fn my_super_function(mut value: usize) {
                value = 10;
            }
        };

        let expected: TokenStream = syn::parse_quote! {
            fn my_super_function(mut value: usize) {
                let timer = ::std::time::Instant::now();

                let mut __micro_timer_inner = | | {
                    value = 10;
                };
                let ret = __micro_timer_inner();

                crate::log::trace!(
                    "Duration of `{}`: {:?}",
                    "my_super_function",
                    timer.elapsed()
                );

                ret
            }
        };
        let output = inner_timed(TokenStream::new(), input);
        assert_eq!(output.to_string(), expected.to_string());
    }

    #[test]
    fn test_ref_mut_output() {
        let input = syn::parse_quote! {
            fn my_super_function(value: &mut String) {
                value.push('a');
            }
        };

        let expected: TokenStream = syn::parse_quote! {
            fn my_super_function(value: &mut String) {
                let timer = ::std::time::Instant::now();

                let mut __micro_timer_inner = | | {
                    value.push('a');
                };
                let ret = __micro_timer_inner();

                crate::log::trace!(
                    "Duration of `{}`: {:?}",
                    "my_super_function",
                    timer.elapsed()
                );

                ret
            }
        };
        let output = inner_timed(TokenStream::new(), input);
        assert_eq!(output.to_string(), expected.to_string());
    }
}