1 use proc_macro::TokenStream;
2 use quote::{quote, ToTokens};
3 use std::collections::HashSet as Set;
4 use syn::fold::{self, Fold};
5 use syn::parse::{Parse, ParseStream, Result};
6 use syn::punctuated::Punctuated;
7 use syn::{parse_macro_input, parse_quote, BinOp, Expr, Ident, ItemFn, Local, Pat, Stmt, Token};
8
9 /// Parses a list of variable names separated by commas.
10 ///
11 /// a, b, c
12 ///
13 /// This is how the compiler passes in arguments to our attribute -- it is
14 /// everything inside the delimiters after the attribute name.
15 ///
16 /// #[trace_var(a, b, c)]
17 /// ^^^^^^^
18 struct Args {
19 vars: Set<Ident>,
20 }
21
22 impl Parse for Args {
parse(input: ParseStream) -> Result<Self>23 fn parse(input: ParseStream) -> Result<Self> {
24 let vars = Punctuated::<Ident, Token![,]>::parse_terminated(input)?;
25 Ok(Args {
26 vars: vars.into_iter().collect(),
27 })
28 }
29 }
30
31 impl Args {
32 /// Determines whether the given `Expr` is a path referring to one of the
33 /// variables we intend to print. Expressions are used as the left-hand side
34 /// of the assignment operator.
should_print_expr(&self, e: &Expr) -> bool35 fn should_print_expr(&self, e: &Expr) -> bool {
36 match *e {
37 Expr::Path(ref e) => {
38 if e.path.leading_colon.is_some() {
39 false
40 } else if e.path.segments.len() != 1 {
41 false
42 } else {
43 let first = e.path.segments.first().unwrap();
44 self.vars.contains(&first.ident) && first.arguments.is_empty()
45 }
46 }
47 _ => false,
48 }
49 }
50
51 /// Determines whether the given `Pat` is an identifier equal to one of the
52 /// variables we intend to print. Patterns are used as the left-hand side of
53 /// a `let` binding.
should_print_pat(&self, p: &Pat) -> bool54 fn should_print_pat(&self, p: &Pat) -> bool {
55 match p {
56 Pat::Ident(ref p) => self.vars.contains(&p.ident),
57 _ => false,
58 }
59 }
60
61 /// Produces an expression that assigns the right-hand side to the left-hand
62 /// side and then prints the value.
63 ///
64 /// // Before
65 /// VAR = INIT
66 ///
67 /// // After
68 /// { VAR = INIT; println!("VAR = {:?}", VAR); }
assign_and_print(&mut self, left: Expr, op: &dyn ToTokens, right: Expr) -> Expr69 fn assign_and_print(&mut self, left: Expr, op: &dyn ToTokens, right: Expr) -> Expr {
70 let right = fold::fold_expr(self, right);
71 parse_quote!({
72 #left #op #right;
73 println!(concat!(stringify!(#left), " = {:?}"), #left);
74 })
75 }
76
77 /// Produces a let-binding that assigns the right-hand side to the left-hand
78 /// side and then prints the value.
79 ///
80 /// // Before
81 /// let VAR = INIT;
82 ///
83 /// // After
84 /// let VAR = { let VAR = INIT; println!("VAR = {:?}", VAR); VAR };
let_and_print(&mut self, local: Local) -> Stmt85 fn let_and_print(&mut self, local: Local) -> Stmt {
86 let Local { pat, init, .. } = local;
87 let init = self.fold_expr(*init.unwrap().expr);
88 let ident = match pat {
89 Pat::Ident(ref p) => &p.ident,
90 _ => unreachable!(),
91 };
92 parse_quote! {
93 let #pat = {
94 #[allow(unused_mut)]
95 let #pat = #init;
96 println!(concat!(stringify!(#ident), " = {:?}"), #ident);
97 #ident
98 };
99 }
100 }
101 }
102
103 /// The `Fold` trait is a way to traverse an owned syntax tree and replace some
104 /// of its nodes.
105 ///
106 /// Syn provides two other syntax tree traversal traits: `Visit` which walks a
107 /// shared borrow of a syntax tree, and `VisitMut` which walks an exclusive
108 /// borrow of a syntax tree and can mutate it in place.
109 ///
110 /// All three traits have a method corresponding to each type of node in Syn's
111 /// syntax tree. All of these methods have default no-op implementations that
112 /// simply recurse on any child nodes. We can override only those methods for
113 /// which we want non-default behavior. In this case the traversal needs to
114 /// transform `Expr` and `Stmt` nodes.
115 impl Fold for Args {
fold_expr(&mut self, e: Expr) -> Expr116 fn fold_expr(&mut self, e: Expr) -> Expr {
117 match e {
118 Expr::Assign(e) => {
119 if self.should_print_expr(&e.left) {
120 self.assign_and_print(*e.left, &e.eq_token, *e.right)
121 } else {
122 Expr::Assign(fold::fold_expr_assign(self, e))
123 }
124 }
125 Expr::Binary(e) if is_assign_op(e.op) => {
126 if self.should_print_expr(&e.left) {
127 self.assign_and_print(*e.left, &e.op, *e.right)
128 } else {
129 Expr::Binary(fold::fold_expr_binary(self, e))
130 }
131 }
132 _ => fold::fold_expr(self, e),
133 }
134 }
135
fold_stmt(&mut self, s: Stmt) -> Stmt136 fn fold_stmt(&mut self, s: Stmt) -> Stmt {
137 match s {
138 Stmt::Local(s) => {
139 if s.init.is_some() && self.should_print_pat(&s.pat) {
140 self.let_and_print(s)
141 } else {
142 Stmt::Local(fold::fold_local(self, s))
143 }
144 }
145 _ => fold::fold_stmt(self, s),
146 }
147 }
148 }
149
is_assign_op(op: BinOp) -> bool150 fn is_assign_op(op: BinOp) -> bool {
151 match op {
152 BinOp::AddAssign(_)
153 | BinOp::SubAssign(_)
154 | BinOp::MulAssign(_)
155 | BinOp::DivAssign(_)
156 | BinOp::RemAssign(_)
157 | BinOp::BitXorAssign(_)
158 | BinOp::BitAndAssign(_)
159 | BinOp::BitOrAssign(_)
160 | BinOp::ShlAssign(_)
161 | BinOp::ShrAssign(_) => true,
162 _ => false,
163 }
164 }
165
166 /// Attribute to print the value of the given variables each time they are
167 /// reassigned.
168 ///
169 /// # Example
170 ///
171 /// ```
172 /// #[trace_var(p, n)]
173 /// fn factorial(mut n: u64) -> u64 {
174 /// let mut p = 1;
175 /// while n > 1 {
176 /// p *= n;
177 /// n -= 1;
178 /// }
179 /// p
180 /// }
181 /// ```
182 #[proc_macro_attribute]
trace_var(args: TokenStream, input: TokenStream) -> TokenStream183 pub fn trace_var(args: TokenStream, input: TokenStream) -> TokenStream {
184 let input = parse_macro_input!(input as ItemFn);
185
186 // Parse the list of variables the user wanted to print.
187 let mut args = parse_macro_input!(args as Args);
188
189 // Use a syntax tree traversal to transform the function body.
190 let output = args.fold_item_fn(input);
191
192 // Hand the resulting function body back to the compiler.
193 TokenStream::from(quote!(#output))
194 }
195