performance_mark_impl/
lib.rs

1//! This crate implements the macro for `performance_mark` and should not be used directly.
2
3use proc_macro2::{TokenStream, TokenTree};
4use quote::ToTokens;
5use syn::{
6    parse2, parse_quote,
7    spanned::Spanned,
8    visit_mut::{visit_stmt_mut, VisitMut},
9    Expr, Item, Stmt,
10};
11
12#[doc(hidden)]
13pub fn performance_mark(attr: TokenStream, item: TokenStream) -> Result<TokenStream, syn::Error> {
14    let mut item = match parse2::<Item>(item.clone()).unwrap() {
15        Item::Fn(function) => function,
16        _ => {
17            return Err(syn::Error::new(
18                item.into_token_stream().span(),
19                "Unexpected token",
20            ))
21        }
22    };
23
24    let mut asyncness = false;
25    let mut log_function = Vec::new();
26
27    for token in attr {
28        match &token {
29            TokenTree::Ident(ident) => {
30                if ident.to_string() == "async" {
31                    asyncness = true;
32                } else {
33                    log_function.push(token);
34                }
35            }
36            TokenTree::Punct(_) => {
37                log_function.push(token);
38            }
39            _ => return Err(syn::Error::new(token.span(), "Unexpected token")),
40        }
41    }
42
43    let log_function = if log_function.is_empty() {
44        None
45    } else {
46        Some(ArbitraryFunction(&log_function))
47    };
48
49    let start_stmt: Stmt = parse_quote!(let start = std::time::Instant::now(););
50
51    item.block.stmts.insert(0, start_stmt);
52
53    let function_name = item.sig.ident.to_string();
54    let mut end_stmts: Vec<Stmt> = parse_quote! {
55        let ctx = performance_mark_attribute::LogContext {
56            function: #function_name.to_string(),
57            duration: std::time::Instant::now().duration_since(start),
58        };
59    };
60
61    if let Some(log_function) = log_function {
62        if asyncness {
63            end_stmts.push(parse_quote! {
64                #log_function(ctx).await;
65            })
66        } else {
67            end_stmts.push(parse_quote! {
68                #log_function(ctx);
69            });
70        }
71    } else {
72        end_stmts.push(parse_quote! {
73            println!("(performance_mark) {} took {:?}", ctx.function, ctx.duration);
74        })
75    }
76
77    let mut visitor = InsertBeforeReturnVisitor {
78        end_stmts: &end_stmts,
79        asyncness,
80    };
81    visitor.visit_item_fn_mut(&mut item);
82    item.block.stmts.extend(end_stmts);
83
84    Ok(item.into_token_stream())
85}
86
87struct InsertBeforeReturnVisitor<'a> {
88    end_stmts: &'a Vec<Stmt>,
89    asyncness: bool,
90}
91
92impl<'a> InsertBeforeReturnVisitor<'a> {
93    fn construct_expr(&self, return_stmt: &Stmt) -> Expr {
94        let stmts = VecStmt(self.end_stmts);
95
96        if self.asyncness {
97            Expr::Await(parse_quote! {
98                async {
99                    #stmts
100                    #return_stmt
101                }.await
102            })
103        } else {
104            Expr::Block(parse_quote! {
105                {
106                    #stmts,
107                    #return_stmt
108                }
109            })
110        }
111    }
112}
113
114impl<'a> VisitMut for InsertBeforeReturnVisitor<'a> {
115    fn visit_stmt_mut(&mut self, stmt: &mut Stmt) {
116        let original_stmt = stmt.clone();
117
118        match stmt {
119            Stmt::Expr(Expr::Return(return_expr), _) => {
120                return_expr
121                    .expr
122                    .replace(Box::new(self.construct_expr(&original_stmt)));
123            }
124            Stmt::Expr(ref mut return_expr, None) => {
125                match return_expr {
126                    Expr::ForLoop(_) | Expr::If(_) | Expr::Loop(_) | Expr::While(_) => {
127                        return visit_stmt_mut(self, stmt);
128                    }
129                    _ => {}
130                }
131
132                *return_expr = self.construct_expr(&original_stmt);
133            }
134            _ => {}
135        }
136    }
137
138    fn visit_expr_closure_mut(&mut self, _: &mut syn::ExprClosure) {
139        // NO-OP, do not visit the inside of closures
140    }
141}
142
143struct VecStmt<'a>(&'a Vec<Stmt>);
144
145impl<'a> ToTokens for VecStmt<'a> {
146    fn to_tokens(&self, tokens: &mut TokenStream) {
147        for stmt in self.0.iter() {
148            stmt.to_tokens(tokens);
149        }
150    }
151}
152
153struct ArbitraryFunction<'a>(&'a Vec<TokenTree>);
154
155impl<'a> ToTokens for ArbitraryFunction<'a> {
156    fn to_tokens(&self, tokens: &mut TokenStream) {
157        for token in self.0.iter() {
158            token.to_tokens(tokens);
159        }
160    }
161}
performance_mark_impl/lib.rs

performance_mark_impl/
lib.rs
