calculagraph/
lib.rs

1//! A handy library for measuring the execution time of function.
2//!
3//! # Use
4//! The use of the crate is through the 4 attribute macros:
5//! `#[timer_println]`, `#[timer_log_trace]`, `#[timer_log_info]`, `#[timer_log_debug]`.
6//!
7//! Just like the name they have, `timer_println` means that it will be using the macro `println!` to
8//! print the results, while `timer_log_trace`, `timer_log_info` and `timer_log_debug` respectively
9//! using macros `log::trace!`, `log::info!` and `log::debug!` to output.
10//!
11//! Of course, when using the last 3 macros, you should link the log facade and its implementation.
12//!
13//! ## Examples
14//! ```toml
15//! [dependencies]
16//! calculagraph = "0.1"
17//! ```
18//! ```rust
19//! use std::{thread, time};
20//! use calculagraph::timer_println;
21//!
22//! #[timer_println(ms)]
23//! fn main() {
24//!     thread::sleep(time::Duration::from_millis(10));
25//!     println!("job done");
26//! }
27//! ```
28//! The above example will print `fn:main cost 10ms` at the end, You can also use the second
29//! argument to define the format string you need.
30//!
31//! ## Note
32//! This macro added two variables that would never conflict in ordinary business, they are
33//! `now_x7bf707c839bc2554fa3f1913a8dc699b68236726c5da18b31f660948ca7f542a267de9b` and
34//! `result_x7bf707c839bc2554fa3f1913a8dc699b68236726c5da18b31f660948ca7f542a267de9b`, just don't
35//! use them intentionally.
36//!
37
38use proc_macro::TokenStream;
39use proc_macro2::TokenStream as TokenStream2;
40use quote::{format_ident, quote};
41use syn::{
42    parse_macro_input, spanned::Spanned, Attribute, AttributeArgs, Block, Error, Item, ItemFn, Lit,
43    Meta, NestedMeta, Path, Result, Signature, Visibility,
44};
45
46// A suffix with a very low probability of conflict,
47// You should not use `now_$UID_SUFFIX$` or `result_$UID_SUFFIX$` as variable name in functions that use the macro.
48// This will cause unexpected result.
49const UID_SUFFIX: &str = "x7bf707c839bc2554fa3f1913a8dc699b68236726c5da18b31f660948ca7f542a267de9b";
50
51/// `std::println!` the execution time after the function is called and executed.
52///
53/// The macro support none, 1 or 2 parameters, [(`TimeUnit`[, `FormatString`])].
54/// The parameter `TimeUnit` supports four types of `s`, `ms`(by default), `us` and `ns`,
55/// The parameter `FormatString` is similar to the format string parameter in `println!`, note that
56/// only one placeholder is supported here, which will fill the time result.
57/// ### Examples
58/// ```
59/// #[timer_println]
60/// fn func() {}
61///
62/// #[timer_println(ns)]
63/// fn func1() {}
64///
65/// #[timer_println(ns, "func2() execution time: {}ns")]
66/// fn func2() {}
67/// ```
68///
69#[proc_macro_attribute]
70pub fn timer_println(attr: TokenStream, input: TokenStream) -> TokenStream {
71    builder(
72        &parse_macro_input!(attr as AttributeArgs),
73        &parse_macro_input!(input as Item),
74        &quote!( println! ),
75    )
76    .unwrap()
77}
78
79/// `log::info!` the execution time after the function is called and executed,
80///
81/// The macro support none, 1 or 2 parameters, [(`TimeUnit`[, `FormatString`])].
82/// The parameter `TimeUnit` supports four types of `s`, `ms`(by default), `us` and `ns`,
83/// The parameter `FormatString` is similar to the format string parameter in `println!`, note that
84/// only one placeholder is supported here, which will fill the time result.
85/// ### Examples
86/// ```
87/// #[timer_log_info]
88/// fn func() {}
89///
90/// #[timer_log_info(ns)]
91/// fn func1() {}
92///
93/// #[timer_log_info(ns, "func2() execution time: {}ns")]
94/// fn func2() {}
95/// ```
96///
97#[proc_macro_attribute]
98pub fn timer_log_info(attr: TokenStream, input: TokenStream) -> TokenStream {
99    builder(
100        &parse_macro_input!(attr as AttributeArgs),
101        &parse_macro_input!(input as Item),
102        &quote!( log::info! ),
103    )
104    .unwrap()
105}
106
107/// `log::debug!` the execution time after the function is called and executed,
108///
109/// The macro support none, 1 or 2 parameters, [(`TimeUnit`[, `FormatString`])].
110/// The parameter `TimeUnit` supports four types of `s`, `ms`(by default), `us` and `ns`,
111/// The parameter `FormatString` is similar to the format string parameter in `println!`, note that
112/// only one placeholder is supported here, which will fill the time result.
113/// ### Examples
114/// ```
115/// #[timer_log_debug]
116/// fn func() {}
117///
118/// #[timer_log_debug(ns)]
119/// fn func1() {}
120///
121/// #[timer_log_debug(ns, "func2() execution time: {}ns")]
122/// fn func2() {}
123/// ```
124///
125#[proc_macro_attribute]
126pub fn timer_log_debug(attr: TokenStream, input: TokenStream) -> TokenStream {
127    builder(
128        &parse_macro_input!(attr as AttributeArgs),
129        &parse_macro_input!(input as Item),
130        &quote!( log::debug! ),
131    )
132    .unwrap()
133}
134
135/// `log::trace!` the execution time after the function is called and executed,
136///
137/// The macro support none, 1 or 2 parameters, [(`TimeUnit`[, `FormatString`])].
138/// The parameter `TimeUnit` supports four types of `s`, `ms`(by default), `us` and `ns`,
139/// The parameter `FormatString` is similar to the format string parameter in `println!`, note that
140/// only one placeholder is supported here, which will fill the time result.
141/// ### Examples
142/// ```
143/// #[timer_log_trace]
144/// fn func() {}
145///
146/// #[timer_log_trace(ns)]
147/// fn func1() {}
148///
149/// #[timer_log_trace(ns, "func2() execution time: {}ns")]
150/// fn func2() {}
151/// ```
152///
153#[proc_macro_attribute]
154pub fn timer_log_trace(attr: TokenStream, input: TokenStream) -> TokenStream {
155    builder(
156        &parse_macro_input!(attr as AttributeArgs),
157        &parse_macro_input!(input as Item),
158        &quote!( log::trace! ),
159    )
160    .unwrap()
161}
162
163#[derive(Debug, Copy, Clone)]
164enum TimeUnit {
165    S,
166    MS,
167    US,
168    NS,
169}
170
171impl TimeUnit {
172    fn quote(self) -> TokenStream2 {
173        match self {
174            TimeUnit::S => quote!(as_secs()),
175            TimeUnit::MS => quote!(as_millis()),
176            TimeUnit::US => quote!(as_micros()),
177            TimeUnit::NS => quote!(as_nanos()),
178        }
179    }
180}
181
182impl std::fmt::Display for TimeUnit {
183    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
184        match *self {
185            TimeUnit::S => write!(f, "s"),
186            TimeUnit::MS => write!(f, "ms"),
187            TimeUnit::US => write!(f, "us"),
188            TimeUnit::NS => write!(f, "ns"),
189        }
190    }
191}
192
193impl std::convert::From<String> for TimeUnit {
194    fn from(s: String) -> Self {
195        match s.to_uppercase().as_str() {
196            "S" => TimeUnit::S,
197            "MS" => TimeUnit::MS,
198            "US" => TimeUnit::US,
199            "NS" => TimeUnit::NS,
200            _ => panic!("Invalid unit of time, only `s`, `ms`, `us`, `ns` are supported"),
201        }
202    }
203}
204
205fn builder(args: &AttributeArgs, body: &Item, outputter: &TokenStream2) -> Result<TokenStream> {
206    let (f_attrs, f_vis, f_sig, f_block) = parse_body(body)?;
207    let f_name = &f_sig.ident.to_string();
208    let (t_unit, o_format) = parse_args(args, f_name)?;
209
210    return Ok((move || {
211        let t_formatter = t_unit.quote();
212        let v_now = format_ident!("now_{}", UID_SUFFIX);
213        let v_result = format_ident!("result_{}", UID_SUFFIX);
214        let f_stmts = &f_block.stmts;
215        quote!(
216            #(#f_attrs)* #f_vis #f_sig {
217                let #v_now = std::time::Instant::now();
218                let #v_result = (move ||{ #(#f_stmts)* })();
219                #outputter(#o_format, #v_now . elapsed(). #t_formatter );
220                return #v_result;
221            }
222        )
223    })()
224    .into());
225
226    // -> (time_unit, output_format_string)
227    fn parse_args(args: &AttributeArgs, fn_name: &String) -> Result<(TimeUnit, String)> {
228        match &args[..] {
229            [] => Ok((TimeUnit::MS, format!("fn:{} cost {{}}ms", fn_name))),
230            [u] => {
231                let u = read_time_unit(u)?;
232                Ok((u, format!("fn:{} cost {{}}{}", fn_name, u)))
233            }
234            [u, f] => Ok((read_time_unit(u)?, read_output_format(f)?)),
235            _ => panic!("Invalid arguments, usage: [(TimeUnit[, OutputFormatString])]"),
236        }
237    }
238
239    fn parse_body(body: &Item) -> Result<(&Vec<Attribute>, &Visibility, &Signature, &Box<Block>)> {
240        match body {
241            Item::Fn(f @ _) => {
242                let ItemFn {
243                    attrs,
244                    vis,
245                    sig,
246                    block,
247                } = f;
248                Ok((attrs, vis, sig, block))
249            }
250            _ => Err(Error::new(
251                body.span(),
252                "Statement other than function are not supported",
253            )),
254        }
255    }
256
257    fn read_time_unit(u: &NestedMeta) -> Result<TimeUnit> {
258        if let NestedMeta::Meta(Meta::Path(Path { segments, .. })) = u {
259            return Ok(segments[0].ident.to_string().into());
260        }
261        Err(syn::Error::new(
262            u.span(),
263            "Invalid argument, `TimeUnit` expected",
264        ))
265    }
266
267    fn read_output_format(f: &NestedMeta) -> Result<String> {
268        match f {
269            NestedMeta::Lit(Lit::Str(s)) => Ok(s.value()),
270            _ => Err(syn::Error::new(
271                f.span(),
272                "Invalid FormatString, the usage is similar with macro `println!` or `format!`",
273            )),
274        }
275    }
276}
calculagraph/lib.rs

calculagraph/
lib.rs
