proc_macro: create a `timeit` decorator

2023-02-19 15:45

Recently, I saw a post on r/rust about creating a macro to check the execution time of a specific function - and my brain just went DECORATOR so... Let's try to build a small proc_macro, shall we?

Main strategy

What we want to achieve is the following: whenever we decorate a function, we will print a message with the elapsed time. As always, the devil is in the details: we want to keep calling the function the same way, so we need to pass the function arguments to the decorated function, as we would do in Python!

We will also leverage the syn and quote crates (as we don't want to manually manage the TokenStream ;) ) and we will simply get the information with the elapsed method of std::time::Instant.

Usage

We want our decorator to be as easy as we can get - so, we just tag the function, and we can read the result back - with the timing information printed on the stdout.

Running this

#[timeit]
fn time_me(something: u8) -> u8 {
    let res = something + 12;
    std::thread::sleep(std::time::Duration::from_millis(4));
    return res;
}
assert_eq!(time_me(30), 42);

#[timeit]
fn time_me_2() {
    std::thread::sleep(std::time::Duration::from_millis(8));
    return;
}
assert_eq!(time_me_2(), ());

will then result in

Took 4 ms to run the function.
Took 8 ms to run the function.

The macro

We are simply wrapping our function, but we want it to be visible with the same name we defined - so, we need to extract the name, and associate it to the wrapper. The wrapped function contains the same body as the original one.

Let's write the whole code and then briefly discuss it

#[proc_macro_attribute]
pub fn timeit(_attrs: TokenStream, item: TokenStream) -> TokenStream {
    let ast: syn::ItemFn = syn::parse(item).expect("this macro accepts functions only");
    let vis = ast.vis;
    let sig = &ast.sig;
    let block = ast.block;
    let sig_inputs = &sig.inputs;
    let output = &sig.output;
    let ident_span = ast.sig.ident.span();
    let fn_ident = syn::Ident::new(
        &format!("__internal{}", ast.sig.ident.to_string()),
        ident_span,
    );
    let sig_inputs_params = sig_inputs
        .iter()
        .map(|arg| {
            if let syn::FnArg::Typed(r) = arg {
                return r.pat.to_owned();
            } else {
                panic!("`self` not expected");
            }
        })
        .collect::<Vec<_>>();

    (quote! {
        #vis #sig {
            #[inline]
            fn #fn_ident(#sig_inputs) #output {
                #block
            }

            let now = std::time::Instant::now();
            let res = #fn_ident(#(#sig_inputs_params,)*);
            let elapsed = now.elapsed();
            println!("Took {} ms to run the function.", elapsed.as_millis());
            return res;
        }
    })
    .into()
}

The fn_ident identifier is the internal name linked to the function - not mandatory, but easier to debug when using cargo expand. sig_inputs_params is a way to extract the parameters, so that I can then give them to the wrapped function. I did not find a way to manage the self parameter yet, but a workaround is to wrap it (double wrap!) like this

struct AStruct {}
impl AStruct {
    pub fn selfmethod(&self) -> u8 {
        std::thread::sleep(std::time::Duration::from_millis(12));
        return 12;
    }
}
let astruct = AStruct {};

#[timeit]
fn wannatimeit(astruct: &AStruct) -> u8 {
    return astruct.selfmethod();
}
assert_eq!(wannatimeit(&astruct), 12);

The rest of the code is trivial: we export the same signature, we run the wrapped function, and we print the result - while returning the result we had.

Final toughts

This (trivial) example can be useful to see how easy it is to run arbitrary code using proc_macros .

Can we do it with declarative macros? Very likely.

Would it be this fun? Doubt so :)