time_graph_macros/
lib.rs

1//! A procedural macro attribute for instrumenting functions with
2//! [`time-graph`].
3//!
4//! [`time-graph`] provides always-on profiling for your code, allowing to
5//! record the execution time of functions, spans inside these functions and the
6//! actual call graph at run-time. This crate provides the
7//! [`#[instrument]`][instrument] procedural macro attribute.
8//!
9//! Note that this macro is also re-exported by the main `time-graph` crate.
10//!
11//!
12//! ## Usage
13//!
14//! First, add this to your `Cargo.toml`:
15//!
16//! ```toml
17//! [dependencies]
18//! time-graph-macros = "0.1.0"
19//! ```
20//!
21//! The [`#[instrument]`][instrument] attribute can now be added to a function
22//! to automatically create a `time-graph` [callsite], and enter the
23//! corresponding [span] when that function is called. For example:
24//!
25//! ```
26//! use time_graph_macros::instrument;
27//!
28//! #[instrument]
29//! pub fn my_function(my_arg: usize) {
30//!     // ...
31//! }
32//!
33//! # fn main() {}
34//! ```
35//!
36//! [`time-graph`]: https://crates.io/crates/time-graph
37//! [instrument]: macro@instrument
38//! [callsite]: https://docs.rs/time-graph/latest/time_graph/struct.CallSite.html
39//! [span]: https://docs.rs/time-graph/latest/time_graph/struct.Span.html
40
41#![allow(clippy::needless_return)]
42
43extern crate proc_macro;
44use proc_macro::TokenStream;
45
46use quote::quote;
47use syn::parse::{Parse, ParseStream};
48use syn::{ItemFn, Signature, LitStr, Token};
49
50
51/// Instruments a function to create and enter a [`time-graph`] [span] every
52/// time the function is called.
53///
54/// # Examples
55/// Instrumenting a function:
56/// ```
57/// # use time_graph_macros::instrument;
58/// #[instrument]
59/// pub fn my_function(my_arg: usize) {
60///     // ...
61/// }
62///
63/// ```
64/// Overriding the generated span's name:
65/// ```
66/// # use time_graph_macros::instrument;
67/// #[instrument(name = "another name")]
68/// pub fn my_function() {
69///     // ...
70/// }
71/// ```
72///
73/// [span]: https://docs.rs/time-graph/latest/time_graph/struct.Span.html
74/// [`time-graph`]: https://github.com/luthaf/time-graph
75#[proc_macro_attribute]
76pub fn instrument(args: TokenStream, tokens: TokenStream) -> TokenStream {
77    let input: ItemFn = syn::parse_macro_input!(tokens as ItemFn);
78    let args: TimedArgs = syn::parse_macro_input!(args as TimedArgs);
79
80    let name = args.name.unwrap_or_else(|| input.sig.ident.to_string());
81
82    let ItemFn {
83        attrs,
84        vis,
85        block,
86        sig,
87        ..
88    } = input;
89
90    let Signature {
91        output: return_type,
92        inputs: params,
93        unsafety,
94        asyncness,
95        constness,
96        abi,
97        ident,
98        generics:
99            syn::Generics {
100                params: gen_params,
101                where_clause,
102                ..
103            },
104        ..
105    } = sig;
106
107    let stream = quote!(
108        #(#attrs) *
109        #vis #constness #unsafety #asyncness #abi fn #ident<#gen_params>(#params) #return_type
110        #where_clause
111        {
112            time_graph::spanned!(#name, {
113                #block
114            })
115        }
116    );
117
118    return stream.into();
119}
120
121struct TimedArgs {
122    name: Option<String>,
123}
124
125mod kw {
126    syn::custom_keyword!(name);
127}
128
129impl Parse for TimedArgs {
130    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
131        let mut args = TimedArgs {
132            name: None,
133        };
134        while !input.is_empty() {
135            let lookahead = input.lookahead1();
136            if lookahead.peek(kw::name) {
137                if args.name.is_some() {
138                    return Err(input.error("expected only a single `name` argument"));
139                }
140                let _ = input.parse::<kw::name>()?;
141                let _ = input.parse::<Token![=]>()?;
142                args.name = Some(input.parse::<LitStr>()?.value());
143            } else if lookahead.peek(LitStr) {
144                if args.name.is_some() {
145                    return Err(input.error("expected only a single `name` argument"));
146                }
147                args.name = Some(input.parse::<LitStr>()?.value());
148            } else {
149                return Err(lookahead.error());
150            }
151        }
152        Ok(args)
153    }
154}
time_graph_macros/lib.rs

time_graph_macros/
lib.rs
