Skip to main content

code_timing_macros/
lib.rs

1#![doc = include_str!("../README.md")]
2#![deny(missing_docs)]
3#![deny(clippy::cargo)]
4#![deny(clippy::pedantic)]
5#![forbid(unsafe_code)]
6#![allow(unreachable_code)]
7
8use proc_macro::TokenStream;
9use quote::quote;
10use syn::ItemFn;
11
12/// Time the duration of a function, either to stdout or via `tracing`.
13#[proc_macro_attribute]
14pub fn time_function(
15    #[allow(unused_variables)] args: TokenStream,
16    input: TokenStream,
17) -> TokenStream {
18    // Do nothing if release and not using the release feature
19    // Do nothing if not testing when using the testing feature
20    #[cfg(any(
21        all(not(debug_assertions), not(feature = "release")),
22        all(not(test), feature = "testing")
23    ))]
24    return input;
25
26    // Parse the input token stream as a function
27    let input = syn::parse_macro_input!(input as ItemFn);
28
29    // Extract the function's signature and body
30    let func_name = &input.sig.ident;
31    let func_block = &input.block;
32    let func_output = &input.sig.output;
33    let func_input = &input.sig.inputs;
34    let func_vis = &input.vis;
35
36    let func_label = if args.is_empty() {
37        format!("{func_name}()")
38    } else {
39        args.to_string()
40    };
41
42    let log_stmt = if cfg!(feature = "tracing") {
43        quote! { ::tracing::trace!("`{}` took {:?}", #func_label, duration); }
44    } else {
45        quote! { println!("`{}` took {:?}", #func_label, duration); }
46    };
47
48    // Generate the wrapped function
49    let output = if input.sig.asyncness.is_some() {
50        quote! {
51            async #func_vis fn #func_name(#func_input) #func_output {
52                let start = ::std::time::Instant::now();
53                let result = (|| async #func_block)().await;
54                let duration: ::std::time::Duration = start.elapsed();
55                #log_stmt
56                result
57            }
58        }
59    } else {
60        quote! {
61            #func_vis fn #func_name(#func_input) #func_output {
62                let start = ::std::time::Instant::now();
63                let result = (|| #func_block)();
64                let duration: ::std::time::Duration = start.elapsed();
65                #log_stmt
66                result
67            }
68        }
69    };
70
71    // Return the generated code as a token stream
72    output.into()
73}
74
75/// Time the duration of code snippet, either to stdout or via `tracing`.
76#[proc_macro]
77pub fn time_snippet(input: TokenStream) -> TokenStream {
78    // Do nothing if release and not using the release feature
79    #[cfg(all(not(debug_assertions), not(feature = "release")))]
80    return input;
81
82    let block: proc_macro2::token_stream::TokenStream = input.into();
83
84    let log_stmt = if cfg!(feature = "tracing") {
85        quote! { ::tracing::trace!("{}:{} took {:?}.", file!(), begin, duration); }
86    } else {
87        quote! { println!("{}:{} took {:?}.", file!(), begin, duration); }
88    };
89
90    let output = quote! {
91        {
92            let begin = line!();
93            let start = ::std::time::Instant::now();
94            let result =
95                #block;
96            let duration: ::std::time::Duration = start.elapsed();
97            #log_stmt
98            result
99        }
100    };
101
102    output.into()
103}