code_timing_macros/
lib.rs

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