timetrace_macros 0.2.2

use proc_macro::TokenStream;

use quote::quote;
use syn::{ItemFn, LitStr, Token, parse_macro_input, punctuated::Punctuated};

#[proc_macro_attribute]
pub fn profile_session_limited(attr: TokenStream, item: TokenStream) -> TokenStream {
    let parser = syn::punctuated::Punctuated::<syn::Expr, Token![,]>::parse_terminated;
    let args = parse_macro_input!(attr with parser);
    let input = parse_macro_input!(item as ItemFn);

    const USAGE: &str =
        "expected: #[profile_session_limited(\"name\", \"trace.json\", max_frames, max_ms)]";

    let mut args_iter = args.iter();

    let session_name = match args_iter.next() {
        Some(syn::Expr::Lit(syn::ExprLit {
            lit: syn::Lit::Str(s),
            ..
        })) => s.clone(),
        _ => {
            return syn::Error::new_spanned(&args, USAGE)
                .to_compile_error()
                .into();
        }
    };

    let filepath = match args_iter.next() {
        Some(syn::Expr::Lit(syn::ExprLit {
            lit: syn::Lit::Str(s),
            ..
        })) => s.clone(),
        _ => {
            return syn::Error::new_spanned(&args, USAGE)
                .to_compile_error()
                .into();
        }
    };

    let max_frames = match args_iter.next() {
        Some(syn::Expr::Lit(syn::ExprLit {
            lit: syn::Lit::Int(n),
            ..
        })) => {
            quote! { Some(#n) }
        }
        None => quote! { None },
        _ => {
            return syn::Error::new_spanned(&args, USAGE)
                .to_compile_error()
                .into();
        }
    };

    let max_ms = match args_iter.next() {
        Some(syn::Expr::Lit(syn::ExprLit {
            lit: syn::Lit::Int(n),
            ..
        })) => {
            quote! { Some(#n) }
        }
        None => quote! { None },
        _ => {
            return syn::Error::new_spanned(&args, USAGE)
                .to_compile_error()
                .into();
        }
    };

    let vis = &input.vis;
    let sig = &input.sig;
    let block = &input.block;

    let expanded = quote! {
        #vis #sig {
            {
                let mut profiler = timetrace::profiler::Profiler::get().lock();

                <timetrace::profiler::Profiler as timetrace::TracingBackend>::begin_session_limited(
                    &mut profiler,
                    #session_name,
                    #filepath,
                    #max_frames,
                    #max_ms,
                );
            }

            let __timetrace_result = (|| #block)();

            {
                let mut profiler = timetrace::profiler::Profiler::get().lock();

                <timetrace::profiler::Profiler as timetrace::TracingBackend>::end_session(
                    &mut profiler,
                ).unwrap();
            }

            __timetrace_result
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_attribute]
pub fn profile_session(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(
        attr with Punctuated::<LitStr, Token![,]>::parse_terminated
    );

    let input = parse_macro_input!(item as ItemFn);

    const USAGE: &str = "expected: #[profile_session(\"name\", \"trace.json\")]";

    let session_name = match args.get(0) {
        Some(s) => s,
        None => {
            return syn::Error::new_spanned(&args, USAGE)
                .to_compile_error()
                .into();
        }
    };

    let filepath = match args.get(1) {
        Some(s) => s,
        None => {
            return syn::Error::new_spanned(&args, USAGE)
                .to_compile_error()
                .into();
        }
    };

    let vis = &input.vis;
    let sig = &input.sig;
    let block = &input.block;

    let expanded = quote! {
        #vis #sig {
            {
                let mut profiler = timetrace::profiler::Profiler::get().lock();

                <timetrace::profiler::Profiler as timetrace::TracingBackend>::begin_session(
                    &mut profiler,
                    #session_name,
                    #filepath,
                );
            }

            let __timetrace_result = (|| #block)();

            {
                let mut profiler = timetrace::profiler::Profiler::get().lock();

                <timetrace::profiler::Profiler as timetrace::TracingBackend>::end_session(
                    &mut profiler,
                ).unwrap();
            }

            __timetrace_result
        }
    };

    TokenStream::from(expanded)
}

#[proc_macro_attribute]
pub fn profile_function(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);

    let vis = &input.vis;
    let sig = &input.sig;
    let block = &input.block;
    let function_name = sig.ident.to_string();

    let expanded = quote! {
        #vis #sig {
            __profile_function!(#function_name);
            #block
        }
    };

    TokenStream::from(expanded)
}