funtime/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::*;
6
7#[proc_macro_attribute]
8pub fn timed(_attrs: TokenStream, item: TokenStream) -> TokenStream {
9    if let Ok(mut fun) = parse::<ItemFn>(item.clone()) {
10        let new_stmts = rewrite_stmts(fun.sig.ident.to_string(), &mut fun.block.stmts);
11        fun.block.stmts = new_stmts;
12        return quote!(#fun).into();
13    }
14
15    if let Ok(mut fun) = parse::<TraitItemMethod>(item.clone()) {
16        if let Some(block) = fun.default.as_mut() {
17            let new_stmts = rewrite_stmts(fun.sig.ident.to_string(), &mut block.stmts);
18            block.stmts = new_stmts;
19            return quote!(#fun).into();
20        }
21    }
22
23    if let Ok(mut fun) = parse::<ImplItemMethod>(item) {
24        let new_stmts = rewrite_stmts(fun.sig.ident.to_string(), &mut fun.block.stmts);
25        fun.block.stmts = new_stmts;
26        return quote!(#fun).into();
27    }
28
29    panic!("`funtime::timed` only works on functions")
30}
31
32fn rewrite_stmts(name: String, stmts: &mut Vec<Stmt>) -> Vec<Stmt> {
33    
34    fn truncate_stmt(stmt: &Stmt, len: usize) -> String {
35        let short =
36            format!("{}", quote::ToTokens::to_token_stream(stmt)).chars().collect::<Vec<_>>();
37
38        let short = if short.len() > len {
39            let mut short = short[..(len - 3)].into_iter().collect::<String>();
40            short.push_str("...");
41            short
42        } else {
43            short.into_iter().collect::<String>()
44        };
45
46        short
47    }
48
49    let setup: Block = parse_quote! {{
50        struct FuntimeTimer {
51            start: std::time::Instant,
52            name: &'static str,
53            buffer: String,
54            prev_mark: Option<std::time::Duration>,
55        }
56
57
58        impl Drop for FuntimeTimer {
59            fn drop(&mut self) {
60                use std::fmt::Write;
61                writeln!(&mut self.buffer, "funtime end: `{}` took {:?}", self.name, self.start.elapsed()).unwrap();
62                print!("{}", &self.buffer);
63            }
64        }
65
66        impl FuntimeTimer {
67            fn new(name: &'static str) -> Self {
68                use std::fmt::Write;
69                let mut buffer = String::new();
70                writeln!(&mut buffer, "funtime start: `{}`", name).unwrap();
71                FuntimeTimer {
72                    start: std::time::Instant::now(),
73                    name,
74                    buffer,
75                    prev_mark: None,
76                }
77            }
78
79            fn mark_elapsed(&mut self, short: &str) {
80                use std::fmt::Write;
81                let mut elapsed = self.start.elapsed();
82                if let Some(prev) = self.prev_mark.replace(elapsed) {
83                    elapsed = elapsed - prev;
84                }
85                writeln!(&mut self.buffer, "  took {:?}: `{}`", elapsed, short).unwrap();
86            }
87        }
88
89        let mut funtime_timer = FuntimeTimer::new(#name);
90
91    }};
92
93    let mut new_stmts = setup.stmts;
94
95    let last = stmts.pop();
96
97    for stmt in stmts.drain(..) {
98        let short = truncate_stmt(&stmt, 40);
99
100        let next_stmt = parse_quote!(funtime_timer.mark_elapsed(#short););
101
102        new_stmts.push(stmt);
103        new_stmts.push(next_stmt);
104    }
105
106    if let Some(stmt) = last {
107        let short = truncate_stmt(&stmt, 40);
108        let new_stmt = parse_quote! {
109            let funtime_return_val = {
110                #stmt
111            };
112        };
113
114        let next_stmt = parse_quote!(funtime_timer.mark_elapsed(#short););
115        let return_stmt = parse_quote!(return funtime_return_val;);
116
117        new_stmts.push(new_stmt);
118        new_stmts.push(next_stmt);
119        new_stmts.push(return_stmt);
120    }
121
122    new_stmts
123}
funtime/lib.rs

funtime/
lib.rs
