rusty_ast/
visitor.rs

1use quote::ToTokens;
2
3pub struct AstVisitor {
4    indent: usize,
5}
6
7impl AstVisitor {
8    pub fn new() -> Self {
9        AstVisitor { indent: 0 }
10    }
11
12    fn print_indent(&self) -> String {
13        " ".repeat(self.indent)
14    }
15}
16
17/// implement Visit trait for AstVisitor
18/// Visit trait is defined in syn::visit
19impl<'ast> syn::visit::Visit<'ast> for AstVisitor {
20    /// visit_item_fn is defined in syn::visit::Visit
21    /// visit_item_fn is called when a Rust function definition is visited
22    ///
23    /// # Arguments
24    /// * `node`: &'ast syn::ItemFn
25    ///
26    /// # Returns
27    /// * `()`
28    fn visit_item_fn(&mut self, node: &'ast syn::ItemFn) {
29        println!("{}Function: {}", self.print_indent(), node.sig.ident);
30        self.indent += 2;
31
32        if !node.sig.inputs.is_empty() {
33            println!("{}Parameters:", self.print_indent());
34            self.indent += 2;
35            for param in &node.sig.inputs {
36                match param {
37                    syn::FnArg::Typed(pat_type) => {
38                        if let syn::Pat::Ident(pat_ident) = &*pat_type.pat {
39                            println!(
40                                "{}Parameter: {} - Type: {}",
41                                self.print_indent(),
42                                pat_ident.ident,
43                                pat_type.ty.to_token_stream()
44                            );
45                        }
46                    }
47                    syn::FnArg::Receiver(receiver) => {
48                        println!(
49                            "{}Self receiver: {}",
50                            self.print_indent(),
51                            receiver.to_token_stream()
52                        );
53                    }
54                }
55            }
56            self.indent -= 2;
57        }
58
59        if let syn::ReturnType::Type(_, return_type) = &node.sig.output {
60            println!(
61                "{}Return type: {}",
62                self.print_indent(),
63                return_type.to_token_stream()
64            );
65        }
66
67        println!("{}Body:", self.print_indent());
68        self.indent += 2;
69        for stmt in &node.block.stmts {
70            self.visit_stmt(stmt);
71        }
72        self.indent -= 4;
73    }
74
75    /// visit_expr is defined in syn::visit::Visit
76    /// visit_expr is called when a Rust expression is visited
77    ///
78    /// # Arguments
79    /// * `node`: &'ast syn::Expr
80    ///
81    /// # Returns
82    /// * `()`
83    fn visit_expr(&mut self, node: &'ast syn::Expr) {
84        match node {
85            syn::Expr::Lit(expr_lit) => match &expr_lit.lit {
86                syn::Lit::Int(lit_int) => {
87                    println!(
88                        "{}Integer literal: {}",
89                        self.print_indent(),
90                        lit_int.base10_digits()
91                    );
92                }
93                syn::Lit::Float(lit_float) => {
94                    println!(
95                        "{}Float literal: {}",
96                        self.print_indent(),
97                        lit_float.base10_digits()
98                    );
99                }
100                syn::Lit::Str(lit_str) => {
101                    println!(
102                        "{}String literal: \"{}\"",
103                        self.print_indent(),
104                        lit_str.value()
105                    );
106                }
107                syn::Lit::Bool(lit_bool) => {
108                    println!("{}Boolean literal: {}", self.print_indent(), lit_bool.value);
109                }
110                _ => {
111                    println!(
112                        "{}Other literal: {}",
113                        self.print_indent(),
114                        expr_lit.to_token_stream()
115                    );
116                }
117            },
118            syn::Expr::Binary(expr_bin) => {
119                let op = match expr_bin.op {
120                    syn::BinOp::Add(_) => "+",
121                    syn::BinOp::Sub(_) => "-",
122                    syn::BinOp::Mul(_) => "*",
123                    syn::BinOp::Div(_) => "/",
124                    syn::BinOp::Eq(_) => "==",
125                    syn::BinOp::Lt(_) => "<",
126                    syn::BinOp::Le(_) => "<=",
127                    syn::BinOp::Ne(_) => "!=",
128                    syn::BinOp::Ge(_) => ">=",
129                    syn::BinOp::Gt(_) => ">",
130                    _ => "other_operator",
131                };
132                println!("{}Binary expression: {}", self.print_indent(), op);
133
134                println!("{}Left:", self.print_indent());
135                self.indent += 2;
136                self.visit_expr(&expr_bin.left);
137                self.indent -= 2;
138
139                println!("{}Right:", self.print_indent());
140                self.indent += 2;
141                self.visit_expr(&expr_bin.right);
142                self.indent -= 2;
143            }
144            syn::Expr::Call(expr_call) => {
145                println!("{}Function call:", self.print_indent());
146
147                println!("{}Function:", self.print_indent());
148                self.indent += 2;
149                self.visit_expr(&expr_call.func);
150                self.indent -= 2;
151
152                if !expr_call.args.is_empty() {
153                    println!("{}Arguments:", self.print_indent());
154                    self.indent += 2;
155                    for arg in &expr_call.args {
156                        self.visit_expr(arg);
157                    }
158                    self.indent -= 2;
159                }
160            }
161            syn::Expr::Path(expr_path) => {
162                println!(
163                    "{}Identifier: {}",
164                    self.print_indent(),
165                    expr_path.to_token_stream()
166                );
167            }
168            syn::Expr::If(expr_if) => {
169                println!("{}If statement:", self.print_indent());
170
171                println!("{}Condition:", self.print_indent());
172                self.indent += 2;
173                self.visit_expr(&expr_if.cond);
174                self.indent -= 2;
175
176                println!("{}Then branch:", self.print_indent());
177                self.indent += 2;
178                for stmt in &expr_if.then_branch.stmts {
179                    self.visit_stmt(stmt);
180                }
181                self.indent -= 2;
182
183                if let Some((_, else_branch)) = &expr_if.else_branch {
184                    println!("{}Else branch:", self.print_indent());
185                    self.indent += 2;
186                    self.visit_expr(&else_branch);
187                    self.indent -= 2;
188                }
189            }
190            syn::Expr::Loop(expr_loop) => {
191                println!("{}Loop:", self.print_indent());
192                self.indent += 2;
193                for stmt in &expr_loop.body.stmts {
194                    self.visit_stmt(stmt);
195                }
196                self.indent -= 2;
197            }
198            syn::Expr::While(expr_while) => {
199                println!("{}While loop:", self.print_indent());
200
201                println!("{}Condition:", self.print_indent());
202                self.indent += 2;
203                self.visit_expr(&expr_while.cond);
204                self.indent -= 2;
205
206                println!("{}Body:", self.print_indent());
207                self.indent += 2;
208                for stmt in &expr_while.body.stmts {
209                    self.visit_stmt(stmt);
210                }
211                self.indent -= 2;
212            }
213            syn::Expr::Return(expr_return) => {
214                println!("{}Return statement:", self.print_indent());
215                if let Some(expr) = &expr_return.expr {
216                    self.indent += 2;
217                    self.visit_expr(expr);
218                    self.indent -= 2;
219                }
220            }
221            _ => {
222                println!(
223                    "{}Other expression: {}",
224                    self.print_indent(),
225                    node.to_token_stream()
226                );
227            }
228        }
229    }
230
231    /// visit_stmt is defined in syn::visit::Visit
232    /// visit_stmt is called when a Rust statement is visited
233    ///
234    /// # Arguments
235    /// * `node`: &'ast syn::Stmt
236    ///
237    /// # Returns
238    /// * `()`
239    fn visit_stmt(&mut self, node: &'ast syn::Stmt) {
240        match node {
241            syn::Stmt::Local(local) => {
242                println!("{}Variable declaration:", self.print_indent());
243                if let syn::Pat::Ident(pat_ident) = &local.pat {
244                    println!("{}Name: {}", self.print_indent(), pat_ident.ident);
245                }
246
247                if let Some(init) = &local.init {
248                    println!("{}Initializer:", self.print_indent());
249                    self.indent += 2;
250                    self.visit_expr(&init.expr);
251                    self.indent -= 2;
252                }
253            }
254
255            syn::Stmt::Expr(expr, _) => {
256                println!("{}Expression statement:", self.print_indent());
257                self.indent += 2;
258                self.visit_expr(expr);
259                self.indent -= 2;
260            }
261
262            syn::Stmt::Item(item) => match item {
263                syn::Item::Fn(item_fn) => {
264                    self.visit_item_fn(item_fn);
265                }
266                syn::Item::Struct(item_struct) => {
267                    println!("{}Struct: {}", self.print_indent(), item_struct.ident);
268                    if !item_struct.fields.is_empty() {
269                        println!("{}Fields:", self.print_indent());
270                        self.indent += 2;
271                        for field in &item_struct.fields {
272                            if let Some(ident) = &field.ident {
273                                println!(
274                                    "{}Field: {} - Type: {}",
275                                    self.print_indent(),
276                                    ident,
277                                    field.ty.to_token_stream()
278                                );
279                            } else {
280                                println!(
281                                    "{}Tuple field: {}",
282                                    self.print_indent(),
283                                    field.ty.to_token_stream()
284                                );
285                            }
286                        }
287                        self.indent -= 2;
288                    }
289                }
290                syn::Item::Enum(item_enum) => {
291                    println!("{}Enum: {}", self.print_indent(), item_enum.ident);
292                    if !item_enum.variants.is_empty() {
293                        println!("{}Variants:", self.print_indent());
294                        self.indent += 2;
295                        for variant in &item_enum.variants {
296                            println!("{}Variant: {}", self.print_indent(), variant.ident);
297                        }
298                        self.indent -= 2;
299                    }
300                }
301                _ => {
302                    println!(
303                        "{}Other item: {}",
304                        self.print_indent(),
305                        item.to_token_stream()
306                    );
307                }
308            },
309            // TODO: add other statement
310            _ => {
311                println!(
312                    "{}Other statement: {}",
313                    self.print_indent(),
314                    node.to_token_stream()
315                );
316            }
317        }
318    }
319}