auto_bench_fct_macros 1.0.0

use proc_macro::TokenStream;
use std::sync::atomic::AtomicU32;
use syn::{parse_macro_input, ItemFn};

/// Add the `#[auto_bench_fct]` macro attribute to a function to enable automatic benchmarking.
/// It will store in a map the number of calls and the total duration of the function execution.
///
/// If the feature `disable` is enabled, the macro will do nothing.
///
/// # Examples
///
/// ```
/// #[auto_bench_fct]
/// fn benchmarked_function() {
///     // Do heavy computation here
/// }
/// ```
#[proc_macro_attribute]
pub fn auto_bench_fct(_attr: TokenStream, item: TokenStream) -> TokenStream {
    #[cfg(feature = "disable")]
    {
        return item;
    }
    bench_fct(_attr, item)
}

/// Add the `#[auto_bench_fct_hy]` macro attribute to a function to enable automatic benchmarking
/// with hierarchy computation. It will do the same as `#[auto_bench_fct]`, but it will also track
/// he call stack, allowing to display the hierarchy of function calls.
///
/// If the feature `disable` is enabled, the macro will do nothing.
/// If the feature `disable_hy` is enabled, the macro will behave like `#[auto_bench_fct]`.
///
/// # Examples
///
/// ```
/// #[auto_bench_fct_hy]
/// fn benchmarked_function() {
///     // Do heavy computation here
/// }
/// ```
#[proc_macro_attribute]
pub fn auto_bench_fct_hy(_attr: TokenStream, item: TokenStream) -> TokenStream {
    #[cfg(feature = "disable")]
    {
        return item;
    }
    #[cfg(feature = "disable_hy")]
    {
        return bench_fct(_attr, item);
    }
    bench_fct_hy(_attr, item)
}

/// Function index counter to ensure unique function identifiers.
/// Essentially counts the number of functions processed by the macro.
static COUNTER: AtomicU32 = AtomicU32::new(1);

/// Internal function to handle the actual macro logic for `auto_bench_fct`.
fn bench_fct(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    let name = &input.sig.ident;
    let block = &input.block;
    let sig = &input.sig;
    let vis = &input.vis;

    let function_key = name.to_string();
    let function_num = COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);

    let output = quote::quote! {
        #vis #sig {
            let start = std::time::Instant::now();
            let result = (|| #block)();
            let duration = start.elapsed();

            let mut metrics = auto_bench_fct::FUNCTION_METRICS.lock().unwrap();
            let entry = metrics.entry((#function_key.to_string(), #function_num)).or_insert((0, std::time::Duration::new(0, 0)));
            entry.0 += 1;
            entry.1 += duration;

            result
        }
    };

    output.into()
}

/// Internal function to handle the actual macro logic for `auto_bench_fct_hy`.
fn bench_fct_hy(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    let name = &input.sig.ident;
    let block = &input.block;
    let sig = &input.sig;
    let vis = &input.vis;

    let function_key = name.to_string();
    let function_num = COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);

    let output = quote::quote! {
        #vis #sig {
            let mut parent = auto_bench_fct::CALL_STACK.with(|stack| stack.borrow().last().cloned().unwrap_or(0));
            let is_recursive = #function_num == parent;
            if !is_recursive {
                auto_bench_fct::CALL_STACK.with(|stack| stack.borrow_mut().push(#function_num));
            }
            let start = std::time::Instant::now();
            let result = (|| #block)();
            let duration = start.elapsed();
            if !is_recursive {
                auto_bench_fct::CALL_STACK.with(|stack| stack.borrow_mut().pop());

                let mut metrics = auto_bench_fct::FUNCTION_METRICS.lock().unwrap();
                let entry = metrics.entry((#function_key.to_string(), #function_num)).or_insert((0, std::time::Duration::new(0, 0)));
                entry.0 += 1;
                entry.1 += duration;

                let mut metrics_hy = auto_bench_fct::FUNCTION_METRICS_HIERARCHY.lock().unwrap();
                let entry = metrics_hy
                    .entry(auto_bench_fct::CALL_STACK.with(|stack| stack.borrow().clone()))
                    .or_insert(std::collections::HashMap::new())
                    .entry((#function_key.to_string(), #function_num))
                    .or_insert((0, std::time::Duration::new(0, 0)));
                entry.0 += 1;
                entry.1 += duration;
            }
            result
        }
    };
    output.into()
}