ast/
visitor.rs

1use std::ops::ControlFlow;
2
3use crate::item::{Item, ItemKind, Module};
4use crate::Block;
5use crate::{
6    expr::ExpressionKind,
7    stmt::{Statement, StatementKind},
8    types::{Type, TypeKind},
9    Expression,
10};
11
12pub trait Visitor {
13    type Result: VisitorResult;
14
15    fn visit_expression(&mut self, expr: &Expression) -> Self::Result {
16        walk_expression(self, expr)
17    }
18
19    fn visit_statement(&mut self, stmt: &Statement) -> Self::Result {
20        walk_statement(self, stmt)
21    }
22
23    fn visit_type(&mut self, ty: &Type) -> Self::Result { walk_type(self, ty) }
24
25    fn visit_module(&mut self, prog: &Module) -> Self::Result { walk_module(self, prog) }
26
27    fn visit_item(&mut self, item: &Item) -> Self::Result {
28        walk_item(self, item)
29    }
30}
31
32pub fn walk_type<V>(v: &mut V, ty: &Type) -> V::Result
33where
34    V: Visitor + ?Sized,
35{
36    match &ty.kind {
37        TypeKind::Array { ty, .. } => v.visit_type(ty),
38        TypeKind::Ref { of, .. } => v.visit_type(of),
39        TypeKind::I8(_) | TypeKind::I16(_) | TypeKind::I32(_) | TypeKind::I64(_) |
40        TypeKind::U8(_) | TypeKind::U16(_) | TypeKind::U32(_) | TypeKind::U64(_) |
41        TypeKind::F32(_) | TypeKind::F64(_) | TypeKind::Bool(_) | TypeKind::Char(_) |
42        TypeKind::Path(_) | TypeKind::Empty(_) => V::Result::output(),
43    }
44}
45
46pub fn walk_item<V>(v: &mut V, item: &Item) -> V::Result
47where
48    V: Visitor + ?Sized,
49{
50    match &item.kind {
51        ItemKind::Variable { ty, init, .. } => {
52            if let Some(ty) = ty {
53                v.visit_type(ty);
54            }
55            if let Some(init) = init {
56                v.visit_expression(init);
57            }
58            V::Result::output()
59        }
60        ItemKind::Function {
61            params: args,
62            return_type,
63            body,
64            ..
65        } => {
66            for param in args {
67                v.visit_type(&param.ty);
68            }
69            if let Some(rty) = return_type {
70                v.visit_type(rty);
71            }
72            if let Some(body) = body {
73                for stmt in &body.val {
74                    v.visit_statement(stmt);
75                }
76            }
77            V::Result::output()
78        }
79        ItemKind::Struct { fields, .. } => {
80            for f in &fields.val {
81                v.visit_type(&f.ty);
82            }
83            V::Result::output()
84        },
85        ItemKind::Mod(m) => v.visit_module(m),
86        ItemKind::Use { .. } => V::Result::output(),
87    }
88}
89
90pub fn walk_statement<V>(v: &mut V, stmt: &Statement) -> V::Result
91where
92    V: Visitor + ?Sized,
93{
94    match &stmt.kind {
95        StatementKind::Expression(expression, _) => v.visit_expression(expression),
96        StatementKind::Item(item) => v.visit_item(item),
97        StatementKind::Block(block) => {
98            for stmt in &block.val {
99                v.visit_statement(stmt);
100            }
101            V::Result::output()
102        }
103        StatementKind::If {
104            cond,
105            if_body,
106            else_body,
107            ..
108        } => {
109            v.visit_expression(&cond.val);
110            v.visit_statement(if_body);
111            if let Some(e) = else_body {
112                v.visit_statement(e);
113            }
114            V::Result::output()
115        }
116        StatementKind::While { cond, body, .. } => {
117            v.visit_expression(cond);
118            v.visit_statement(body);
119            V::Result::output()
120        }
121        StatementKind::For {
122            init,
123            cond,
124            inc,
125            body,
126            ..
127        } => {
128            if let Some(init) = init {
129                v.visit_item(init);
130            }
131            if let Some(cond) = cond {
132                v.visit_expression(cond);
133            }
134            if let Some(inc) = inc {
135                v.visit_expression(inc);
136            }
137
138            v.visit_statement(body);
139
140            V::Result::output()
141        }
142        StatementKind::Empty(_) | StatementKind::Break(_) | StatementKind::Continue(_) => {
143            V::Result::output()
144        }
145        StatementKind::Return { expr, .. } => {
146            if let Some(expr) = expr {
147                v.visit_expression(expr);
148            }
149            V::Result::output()
150        }
151    }
152}
153
154pub fn walk_expression<V>(v: &mut V, expr: &Expression) -> V::Result
155where
156    V: Visitor + ?Sized,
157{
158    match &expr.kind {
159        ExpressionKind::Unary { expr, .. } => v.visit_expression(expr),
160        ExpressionKind::Paren(pexpr) => v.visit_expression(&pexpr.val),
161        ExpressionKind::Cast { expr, ty, .. } => {
162            v.visit_expression(expr);
163            v.visit_type(ty)
164        }
165        ExpressionKind::Binary { left, right, .. } => {
166            v.visit_expression(left);
167            v.visit_expression(right);
168            V::Result::output()
169        }
170        ExpressionKind::Ternary {
171            cond,
172            if_true,
173            if_false,
174        } => {
175            v.visit_expression(cond);
176            v.visit_expression(if_true);
177            v.visit_expression(if_false);
178            V::Result::output()
179        }
180        ExpressionKind::Path(_) | ExpressionKind::Literal(_) => V::Result::output(),
181        ExpressionKind::Call { callee, args } => {
182            v.visit_expression(callee);
183            for a in &args.val {
184                v.visit_expression(a);
185            }
186            V::Result::output()
187        }
188        ExpressionKind::ArrayAccess { arr, index, .. } => {
189            v.visit_expression(arr);
190            v.visit_expression(index);
191            V::Result::output()
192        }
193        ExpressionKind::StructAccess { st, .. } => v.visit_expression(st),
194    }
195}
196
197pub fn walk_module<V>(v: &mut V, module: &Module) -> V::Result
198where
199    V: Visitor + ?Sized,
200{
201    use crate::ModuleBody;
202
203    match &module.body {
204        ModuleBody::Inline(Block { val: items, ..}) |
205        ModuleBody::Extern { items, .. } |
206        ModuleBody::Slf(items, _) => {
207            for item in items {
208                v.visit_item(item);
209            }
210        }
211    }
212    V::Result::output()
213}
214
215pub trait VisitorResult {
216    type T;
217    type Residual;
218
219    fn output() -> Self;
220    fn from_residual(residual: Self::Residual) -> Self;
221    fn from_branch(b: ControlFlow<Self::Residual, Self::T>) -> Self;
222    fn branch(self) -> ControlFlow<Self::Residual, Self::T>;
223}
224
225impl VisitorResult for () {
226    type T = ();
227    type Residual = core::convert::Infallible;
228
229    fn output() -> Self {}
230    fn from_residual(_residual: Self::Residual) -> Self {}
231    fn from_branch(_b: ControlFlow<Self::Residual>) -> Self {}
232    fn branch(self) -> ControlFlow<Self::Residual> { ControlFlow::Continue(()) }
233}
234
235impl<B, C: Default> VisitorResult for ControlFlow<B, C> {
236    type T = C;
237    type Residual = B;
238
239    fn output() -> Self { ControlFlow::Continue(Self::T::default()) }
240
241    fn from_residual(residual: Self::Residual) -> Self { ControlFlow::Break(residual) }
242
243    fn from_branch(b: ControlFlow<Self::Residual, Self::T>) -> Self { b }
244
245    fn branch(self) -> ControlFlow<Self::Residual, Self::T> { self }
246}
247
248impl<T: Default, E> VisitorResult for Result<T, E> {
249    type T = T;
250    type Residual = E;
251
252    fn output() -> Self { Result::Ok(T::default()) }
253
254    fn from_residual(residual: Self::Residual) -> Self { Self::Err(residual) }
255
256    fn from_branch(b: ControlFlow<Self::Residual, Self::T>) -> Self {
257        match b {
258            ControlFlow::Continue(c) => Ok(c),
259            ControlFlow::Break(b) => Err(b),
260        }
261    }
262
263    fn branch(self) -> ControlFlow<Self::Residual, Self::T> {
264        match self {
265            Ok(t) => ControlFlow::Continue(t),
266            Err(err) => ControlFlow::Break(err),
267        }
268    }
269}
270
271impl<T> VisitorResult for Option<T> {
272    type T = T;
273    type Residual = ();
274
275    fn output() -> Self { None }
276
277    fn from_residual(_residual: Self::Residual) -> Self { None }
278
279    fn from_branch(b: ControlFlow<Self::Residual, Self::T>) -> Self {
280        match b {
281            ControlFlow::Continue(c) => Some(c),
282            ControlFlow::Break(()) => None,
283        }
284    }
285
286    fn branch(self) -> ControlFlow<Self::Residual, Self::T> {
287        match self {
288            Some(e) => ControlFlow::Continue(e),
289            None => ControlFlow::Break(()),
290        }
291    }
292}
ast/visitor.rs

ast/
visitor.rs
