fn_util/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, ItemFn};
4
5/// Simply record the execution time of a function
6#[proc_macro_attribute]
7pub fn bench(attr: TokenStream, item: TokenStream) -> TokenStream {
8    let ItemFn { attrs, vis, sig, block } = parse_macro_input!(item as ItemFn);
9
10    let name = match syn::parse::<syn::Ident>(attr) {
11        Ok(x) => x.to_string(),
12        Err(_) => sig.ident.to_string(),
13    };
14
15    let ts = quote! {
16        #(#attrs)* #vis #sig {
17            let __t = std::time::Instant::now();
18            let __result = #block;
19            match std::env::var(crate::BENCH_FLAG) {
20                Ok(_) => {
21                    println!("{}: {:?}", #name, __t.elapsed());
22                },
23                Err(_) => {}
24            };
25            __result
26        }
27    };
28    ts.into()
29}
30
31/// Simply pause the function until any key is pressed
32#[proc_macro_attribute]
33pub fn pause(attr: TokenStream, item: TokenStream) -> TokenStream {
34    let ItemFn { attrs, vis, sig, block } = parse_macro_input!(item as ItemFn);
35    let fn_name = sig.ident.to_string();
36
37    let token_flush_pause = quote! {
38        {
39            use std::io::Write;
40            std::io::stdout().flush().unwrap();
41        }
42        std::io::stdin().read_line(&mut String::new()).unwrap();
43    };
44    let pause_print = quote! {
45        print!("Press enter to start '{}'... ", #fn_name);
46        #token_flush_pause
47    };
48
49    let pause_end = match syn::parse::<syn::Ident>(attr) {
50        Ok(x) => {
51            if &x.to_string() == "both" {
52                quote! {
53                    print!("Press enter to end '{}'... ", #fn_name);
54                    #token_flush_pause
55                }
56            } else {
57                quote! {}
58            }
59        }
60        Err(_) => quote! {},
61    };
62
63    let ts = quote! {
64        #(#attrs)* #vis #sig {
65            #pause_print;
66            let __result = #block;
67            #pause_end;
68            __result
69        }
70    };
71    ts.into()
72}
fn_util/lib.rs

fn_util/
lib.rs
