mest_core/
visitor.rs

1use crate::ast::{BinOp, Expr, ExprKind, Literal, Pat, UnaryOp};
2use lasso::Rodeo;
3
4pub trait Visitor<'bump, Ctx>: Sized {
5    fn visit_expr(&mut self, _expr: &'bump Expr<'bump>, _ctx: &mut Ctx) {}
6
7    fn visit_pat(&mut self, _pat: &'bump Pat<'bump>, _ctx: &mut Ctx) {}
8
9    fn walk_expr(&mut self, expr: &'bump Expr<'bump>, ctx: &mut Ctx) {
10        self.visit_expr(expr, ctx);
11        match &*expr.kind {
12            ExprKind::Literal(_) | ExprKind::Ident(_) => {}
13            ExprKind::If {
14                cond,
15                then_expr,
16                else_expr,
17            } => {
18                self.walk_expr(cond, ctx);
19                self.walk_expr(then_expr, ctx);
20                self.walk_expr(else_expr, ctx);
21            }
22            ExprKind::BinOp { op: _, lhs, rhs } => {
23                self.walk_expr(lhs, ctx);
24                self.walk_expr(rhs, ctx);
25            }
26            ExprKind::UnaryOp { op: _, rhs } => {
27                self.walk_expr(rhs, ctx);
28            }
29            ExprKind::Let {
30                name: _,
31                value,
32                body,
33                rec: _,
34            } => {
35                self.walk_expr(value, ctx);
36                self.walk_expr(body, ctx);
37            }
38            ExprKind::Match { scrutinee, arms } => {
39                self.walk_expr(scrutinee, ctx);
40                for (pat, body) in arms.iter() {
41                    self.walk_pat(pat, ctx);
42                    self.walk_expr(body, ctx);
43                }
44            }
45            ExprKind::Abs { param: _, body } => {
46                self.walk_expr(body, ctx);
47            }
48            ExprKind::App { func, arg } => {
49                self.walk_expr(func, ctx);
50                self.walk_expr(arg, ctx);
51            }
52        }
53    }
54
55    fn walk_pat(&mut self, pat: &'bump Pat<'bump>, ctx: &mut Ctx) {
56        self.visit_pat(pat, ctx);
57        match pat {
58            Pat::Wildcard | Pat::Var(_) | Pat::Lit(_) => {}
59            Pat::Or(a, b) => {
60                self.walk_pat(a, ctx);
61                self.walk_pat(b, ctx);
62            }
63        }
64    }
65}
66
67// TODO: VisitMut requires `kind: &'bump mut ExprKind` on `Expr` to walk children.
68// Add once the AST supports mutable node access.
69
70pub struct PrintCtx<'a> {
71    pub rodeo: &'a Rodeo,
72    pub indent: usize,
73}
74
75impl<'a> PrintCtx<'a> {
76    pub fn new(rodeo: &'a Rodeo) -> Self {
77        Self { rodeo, indent: 0 }
78    }
79}
80
81pub struct AstPrinter;
82
83impl AstPrinter {
84    pub fn print_expr<'bump>(expr: &'bump Expr<'bump>, ctx: &mut PrintCtx) {
85        AstPrinter.visit_expr(expr, ctx);
86    }
87}
88
89impl<'bump> Visitor<'bump, PrintCtx<'_>> for AstPrinter {
90    fn visit_expr(&mut self, expr: &'bump Expr<'bump>, ctx: &mut PrintCtx<'_>) {
91        let indent = "  ".repeat(ctx.indent);
92        match &*expr.kind {
93            ExprKind::Literal(lit) => match lit {
94                Literal::Int(n) => println!("{indent}Int({n})"),
95                Literal::Float(f) => println!("{indent}Float({f})"),
96                Literal::Bool(b) => println!("{indent}Bool({b})"),
97            },
98            ExprKind::Ident(ident) => {
99                println!("{indent}Ident({})", ctx.rodeo.resolve(&ident.0));
100            }
101            ExprKind::If {
102                cond,
103                then_expr,
104                else_expr,
105            } => {
106                println!("{indent}If");
107                ctx.indent += 1;
108                self.visit_expr(cond, ctx);
109                println!("{}Then", "  ".repeat(ctx.indent - 1));
110                self.visit_expr(then_expr, ctx);
111                println!("{}Else", "  ".repeat(ctx.indent - 1));
112                self.visit_expr(else_expr, ctx);
113                ctx.indent -= 1;
114            }
115            ExprKind::BinOp { op, lhs, rhs } => {
116                let op_str = match op {
117                    BinOp::Eq => "==",
118                    BinOp::NotEq => "!=",
119                    BinOp::Lt => "<",
120                    BinOp::Gt => ">",
121                    BinOp::Le => "<=",
122                    BinOp::Ge => ">=",
123                    BinOp::And => "&&",
124                    BinOp::Or => "||",
125                    BinOp::Add => "+",
126                    BinOp::Sub => "-",
127                    BinOp::Mul => "*",
128                    BinOp::Div => "/",
129                    BinOp::Pow => "^",
130                };
131                println!("{indent}BinOp({op_str})");
132                ctx.indent += 1;
133                self.visit_expr(lhs, ctx);
134                self.visit_expr(rhs, ctx);
135                ctx.indent -= 1;
136            }
137            ExprKind::UnaryOp { op, rhs } => {
138                let op_str = match op {
139                    UnaryOp::Neg => "-",
140                    UnaryOp::Not => "!",
141                };
142                println!("{indent}UnaryOp({op_str})");
143                ctx.indent += 1;
144                self.visit_expr(rhs, ctx);
145                ctx.indent -= 1;
146            }
147            ExprKind::Let {
148                name,
149                value,
150                body,
151                rec,
152            } => {
153                if *rec {
154                    println!("{indent}LetRec({})", ctx.rodeo.resolve(&name.0));
155                } else {
156                    println!("{indent}Let({})", ctx.rodeo.resolve(&name.0));
157                }
158                ctx.indent += 1;
159                self.visit_expr(value, ctx);
160                self.visit_expr(body, ctx);
161                ctx.indent -= 1;
162            }
163            ExprKind::Match { scrutinee, arms } => {
164                println!("{indent}Match");
165                ctx.indent += 1;
166                self.visit_expr(scrutinee, ctx);
167                for (pat, body) in arms.iter() {
168                    self.visit_pat(pat, ctx);
169                    self.visit_expr(body, ctx);
170                }
171                ctx.indent -= 1;
172            }
173            ExprKind::Abs { param, body } => {
174                println!("{indent}Lambda({})", ctx.rodeo.resolve(&param.0));
175                ctx.indent += 1;
176                self.visit_expr(body, ctx);
177                ctx.indent -= 1;
178            }
179            ExprKind::App { func, arg } => {
180                println!("{indent}App");
181                ctx.indent += 1;
182                self.visit_expr(func, ctx);
183                self.visit_expr(arg, ctx);
184                ctx.indent -= 1;
185            }
186        }
187    }
188
189    fn visit_pat(&mut self, pat: &'bump Pat<'bump>, ctx: &mut PrintCtx<'_>) {
190        let indent = "  ".repeat(ctx.indent);
191        match pat {
192            Pat::Wildcard => println!("{indent}Wildcard"),
193            Pat::Var(ident) => {
194                println!("{indent}Var({})", ctx.rodeo.resolve(&ident.0));
195            }
196            Pat::Lit(lit) => match lit {
197                Literal::Int(n) => println!("{indent}Lit(Int({n}))"),
198                Literal::Float(f) => println!("{indent}Lit(Float({f}))"),
199                Literal::Bool(b) => println!("{indent}Lit(Bool({b}))"),
200            },
201            Pat::Or(a, b) => {
202                println!("{indent}Or");
203                ctx.indent += 1;
204                self.visit_pat(a, ctx);
205                self.visit_pat(b, ctx);
206                ctx.indent -= 1;
207            }
208        }
209    }
210}
mest_core/visitor.rs

mest_core/
visitor.rs
