use proc_macro::TokenStream;
use quote::quote;
use syn::parse_macro_input;
#[proc_macro_attribute]
pub fn benchmark(_attrs: TokenStream, item: TokenStream) -> TokenStream {
let input_fn = parse_macro_input!(item as syn::ItemFn);
let fn_name = &input_fn.sig.ident;
let fn_block = &input_fn.block;
let func_return = &input_fn.sig.output;
let func_args = &input_fn.sig.inputs;
let expanded = quote! {
fn #fn_name(#func_args) #func_return {
use jemalloc_ctl::{epoch, stats::allocated};
println!("Entering function: {}", stringify!(#fn_name));
let e = epoch::mib().unwrap();
let allocated_before = allocated::mib().unwrap();
e.advance().unwrap();
let memory_before = allocated_before.read().unwrap();
let start = std::time::Instant::now();
let result = (|| #fn_block)();
let duration = start.elapsed().as_millis();
e.advance().unwrap();
let memory_after = allocated_before.read().unwrap();
let memory_used = memory_after - memory_before;
println!("Exiting function: {} (took {} ms) memory used: {:?} bytes", stringify!(#fn_name), duration, memory_used);
result
}
};
TokenStream::from(expanded)
}