leo_ast/passes/
visitor.rs

1// Copyright (C) 2019-2025 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17//! This module contains Visitor trait implementations for the AST.
18//! It implements default methods for each node to be made
19//! given the type of node its visiting.
20
21use crate::*;
22
23// TODO: The Visitor and Reconstructor patterns need a redesign so that the default implementation can easily be invoked though its implemented in an overriding trait.
24// Here is a pattern that seems to work
25// trait ProgramVisitor {
26//     // The trait method that can be overridden
27//     fn visit_program_scope(&mut self);
28//
29//     // Private helper function containing the default implementation
30//     fn default_visit_program_scope(&mut self) {
31//         println!("Do default stuff");
32//     }
33// }
34//
35// struct YourStruct;
36//
37// impl ProgramVisitor for YourStruct {
38//     fn visit_program_scope(&mut self) {
39//         println!("Do custom stuff.");
40//         // Call the default implementation
41//         self.default_visit_program_scope();
42//     }
43// }
44
45/// A Visitor trait for types in the AST.
46pub trait AstVisitor {
47    /* Types */
48    fn visit_type(&mut self, input: &Type) {
49        match input {
50            Type::Array(array_type) => self.visit_array_type(array_type),
51            Type::Composite(composite_type) => self.visit_composite_type(composite_type),
52            Type::Future(future_type) => self.visit_future_type(future_type),
53            Type::Mapping(mapping_type) => self.visit_mapping_type(mapping_type),
54            Type::Tuple(tuple_type) => self.visit_tuple_type(tuple_type),
55            Type::Address
56            | Type::Boolean
57            | Type::Field
58            | Type::Group
59            | Type::Identifier(_)
60            | Type::Integer(_)
61            | Type::Scalar
62            | Type::Signature
63            | Type::String
64            | Type::Numeric
65            | Type::Unit
66            | Type::Err => {}
67        }
68    }
69
70    fn visit_array_type(&mut self, input: &ArrayType) {
71        self.visit_type(&input.element_type);
72        self.visit_expression(&input.length, &Default::default());
73    }
74
75    fn visit_composite_type(&mut self, input: &CompositeType) {
76        input.const_arguments.iter().for_each(|expr| {
77            self.visit_expression(expr, &Default::default());
78        });
79    }
80
81    fn visit_future_type(&mut self, input: &FutureType) {
82        input.inputs.iter().for_each(|input| self.visit_type(input));
83    }
84
85    fn visit_mapping_type(&mut self, input: &MappingType) {
86        self.visit_type(&input.key);
87        self.visit_type(&input.value);
88    }
89
90    fn visit_tuple_type(&mut self, input: &TupleType) {
91        input.elements().iter().for_each(|input| self.visit_type(input));
92    }
93
94    /* Expressions */
95    type AdditionalInput: Default;
96    type Output: Default;
97
98    fn visit_expression(&mut self, input: &Expression, additional: &Self::AdditionalInput) -> Self::Output {
99        match input {
100            Expression::Array(array) => self.visit_array(array, additional),
101            Expression::ArrayAccess(access) => self.visit_array_access(access, additional),
102            Expression::AssociatedConstant(constant) => self.visit_associated_constant(constant, additional),
103            Expression::AssociatedFunction(function) => self.visit_associated_function(function, additional),
104            Expression::Async(async_) => self.visit_async(async_, additional),
105            Expression::Binary(binary) => self.visit_binary(binary, additional),
106            Expression::Call(call) => self.visit_call(call, additional),
107            Expression::Cast(cast) => self.visit_cast(cast, additional),
108            Expression::Struct(struct_) => self.visit_struct_init(struct_, additional),
109            Expression::Err(err) => self.visit_err(err, additional),
110            Expression::Path(path) => self.visit_path(path, additional),
111            Expression::Literal(literal) => self.visit_literal(literal, additional),
112            Expression::Locator(locator) => self.visit_locator(locator, additional),
113            Expression::MemberAccess(access) => self.visit_member_access(access, additional),
114            Expression::Repeat(repeat) => self.visit_repeat(repeat, additional),
115            Expression::Ternary(ternary) => self.visit_ternary(ternary, additional),
116            Expression::Tuple(tuple) => self.visit_tuple(tuple, additional),
117            Expression::TupleAccess(access) => self.visit_tuple_access(access, additional),
118            Expression::Unary(unary) => self.visit_unary(unary, additional),
119            Expression::Unit(unit) => self.visit_unit(unit, additional),
120        }
121    }
122
123    fn visit_array_access(&mut self, input: &ArrayAccess, additional: &Self::AdditionalInput) -> Self::Output {
124        self.visit_expression(&input.array, additional);
125        self.visit_expression(&input.index, additional);
126        Default::default()
127    }
128
129    fn visit_member_access(&mut self, input: &MemberAccess, additional: &Self::AdditionalInput) -> Self::Output {
130        self.visit_expression(&input.inner, additional);
131        Default::default()
132    }
133
134    fn visit_tuple_access(&mut self, input: &TupleAccess, additional: &Self::AdditionalInput) -> Self::Output {
135        self.visit_expression(&input.tuple, additional);
136        Default::default()
137    }
138
139    fn visit_array(&mut self, input: &ArrayExpression, additional: &Self::AdditionalInput) -> Self::Output {
140        input.elements.iter().for_each(|expr| {
141            self.visit_expression(expr, additional);
142        });
143        Default::default()
144    }
145
146    fn visit_associated_constant(
147        &mut self,
148        _input: &AssociatedConstantExpression,
149        _additional: &Self::AdditionalInput,
150    ) -> Self::Output {
151        Default::default()
152    }
153
154    fn visit_associated_function(
155        &mut self,
156        input: &AssociatedFunctionExpression,
157        _additional: &Self::AdditionalInput,
158    ) -> Self::Output {
159        input.arguments.iter().for_each(|arg| {
160            self.visit_expression(arg, &Default::default());
161        });
162        Default::default()
163    }
164
165    fn visit_async(&mut self, input: &AsyncExpression, _additional: &Self::AdditionalInput) -> Self::Output {
166        self.visit_block(&input.block);
167        Default::default()
168    }
169
170    fn visit_binary(&mut self, input: &BinaryExpression, additional: &Self::AdditionalInput) -> Self::Output {
171        self.visit_expression(&input.left, additional);
172        self.visit_expression(&input.right, additional);
173        Default::default()
174    }
175
176    fn visit_call(&mut self, input: &CallExpression, additional: &Self::AdditionalInput) -> Self::Output {
177        input.const_arguments.iter().for_each(|expr| {
178            self.visit_expression(expr, additional);
179        });
180        input.arguments.iter().for_each(|expr| {
181            self.visit_expression(expr, additional);
182        });
183        Default::default()
184    }
185
186    fn visit_cast(&mut self, input: &CastExpression, additional: &Self::AdditionalInput) -> Self::Output {
187        self.visit_expression(&input.expression, additional);
188        Default::default()
189    }
190
191    fn visit_struct_init(&mut self, input: &StructExpression, additional: &Self::AdditionalInput) -> Self::Output {
192        input.const_arguments.iter().for_each(|expr| {
193            self.visit_expression(expr, additional);
194        });
195        for StructVariableInitializer { expression, .. } in input.members.iter() {
196            if let Some(expression) = expression {
197                self.visit_expression(expression, additional);
198            }
199        }
200        Default::default()
201    }
202
203    fn visit_err(&mut self, _input: &ErrExpression, _additional: &Self::AdditionalInput) -> Self::Output {
204        panic!("`ErrExpression`s should not be in the AST at this phase of compilation.")
205    }
206
207    fn visit_path(&mut self, _input: &Path, _additional: &Self::AdditionalInput) -> Self::Output {
208        Default::default()
209    }
210
211    fn visit_literal(&mut self, _input: &Literal, _additional: &Self::AdditionalInput) -> Self::Output {
212        Default::default()
213    }
214
215    fn visit_locator(&mut self, _input: &LocatorExpression, _additional: &Self::AdditionalInput) -> Self::Output {
216        Default::default()
217    }
218
219    fn visit_repeat(&mut self, input: &RepeatExpression, additional: &Self::AdditionalInput) -> Self::Output {
220        self.visit_expression(&input.expr, additional);
221        self.visit_expression(&input.count, additional);
222        Default::default()
223    }
224
225    fn visit_ternary(&mut self, input: &TernaryExpression, additional: &Self::AdditionalInput) -> Self::Output {
226        self.visit_expression(&input.condition, additional);
227        self.visit_expression(&input.if_true, additional);
228        self.visit_expression(&input.if_false, additional);
229        Default::default()
230    }
231
232    fn visit_tuple(&mut self, input: &TupleExpression, additional: &Self::AdditionalInput) -> Self::Output {
233        input.elements.iter().for_each(|expr| {
234            self.visit_expression(expr, additional);
235        });
236        Default::default()
237    }
238
239    fn visit_unary(&mut self, input: &UnaryExpression, additional: &Self::AdditionalInput) -> Self::Output {
240        self.visit_expression(&input.receiver, additional);
241        Default::default()
242    }
243
244    fn visit_unit(&mut self, _input: &UnitExpression, _additional: &Self::AdditionalInput) -> Self::Output {
245        Default::default()
246    }
247
248    /* Statements */
249    fn visit_statement(&mut self, input: &Statement) {
250        match input {
251            Statement::Assert(stmt) => self.visit_assert(stmt),
252            Statement::Assign(stmt) => self.visit_assign(stmt),
253            Statement::Block(stmt) => self.visit_block(stmt),
254            Statement::Conditional(stmt) => self.visit_conditional(stmt),
255            Statement::Const(stmt) => self.visit_const(stmt),
256            Statement::Definition(stmt) => self.visit_definition(stmt),
257            Statement::Expression(stmt) => self.visit_expression_statement(stmt),
258            Statement::Iteration(stmt) => self.visit_iteration(stmt),
259            Statement::Return(stmt) => self.visit_return(stmt),
260        }
261    }
262
263    fn visit_assert(&mut self, input: &AssertStatement) {
264        match &input.variant {
265            AssertVariant::Assert(expr) => self.visit_expression(expr, &Default::default()),
266            AssertVariant::AssertEq(left, right) | AssertVariant::AssertNeq(left, right) => {
267                self.visit_expression(left, &Default::default());
268                self.visit_expression(right, &Default::default())
269            }
270        };
271    }
272
273    fn visit_assign(&mut self, input: &AssignStatement) {
274        self.visit_expression(&input.place, &Default::default());
275        self.visit_expression(&input.value, &Default::default());
276    }
277
278    fn visit_block(&mut self, input: &Block) {
279        input.statements.iter().for_each(|stmt| self.visit_statement(stmt));
280    }
281
282    fn visit_conditional(&mut self, input: &ConditionalStatement) {
283        self.visit_expression(&input.condition, &Default::default());
284        self.visit_block(&input.then);
285        if let Some(stmt) = input.otherwise.as_ref() {
286            self.visit_statement(stmt);
287        }
288    }
289
290    fn visit_const(&mut self, input: &ConstDeclaration) {
291        self.visit_type(&input.type_);
292        self.visit_expression(&input.value, &Default::default());
293    }
294
295    fn visit_definition(&mut self, input: &DefinitionStatement) {
296        if let Some(ty) = input.type_.as_ref() {
297            self.visit_type(ty)
298        }
299        self.visit_expression(&input.value, &Default::default());
300    }
301
302    fn visit_expression_statement(&mut self, input: &ExpressionStatement) {
303        self.visit_expression(&input.expression, &Default::default());
304    }
305
306    fn visit_iteration(&mut self, input: &IterationStatement) {
307        if let Some(ty) = input.type_.as_ref() {
308            self.visit_type(ty)
309        }
310        self.visit_expression(&input.start, &Default::default());
311        self.visit_expression(&input.stop, &Default::default());
312        self.visit_block(&input.block);
313    }
314
315    fn visit_return(&mut self, input: &ReturnStatement) {
316        self.visit_expression(&input.expression, &Default::default());
317    }
318}
319
320/// A Visitor trait for the program represented by the AST.
321pub trait ProgramVisitor: AstVisitor {
322    fn visit_program(&mut self, input: &Program) {
323        input.modules.values().for_each(|module| self.visit_module(module));
324        input.imports.values().for_each(|import| self.visit_import(&import.0));
325        input.stubs.values().for_each(|stub| self.visit_stub(stub));
326        input.program_scopes.values().for_each(|scope| self.visit_program_scope(scope));
327    }
328
329    fn visit_program_scope(&mut self, input: &ProgramScope) {
330        input.consts.iter().for_each(|(_, c)| (self.visit_const(c)));
331        input.structs.iter().for_each(|(_, c)| (self.visit_struct(c)));
332        input.mappings.iter().for_each(|(_, c)| (self.visit_mapping(c)));
333        input.functions.iter().for_each(|(_, c)| (self.visit_function(c)));
334        if let Some(c) = input.constructor.as_ref() {
335            self.visit_constructor(c);
336        }
337    }
338
339    fn visit_module(&mut self, input: &Module) {
340        input.consts.iter().for_each(|(_, c)| (self.visit_const(c)));
341        input.structs.iter().for_each(|(_, c)| (self.visit_struct(c)));
342        input.functions.iter().for_each(|(_, c)| (self.visit_function(c)));
343    }
344
345    fn visit_stub(&mut self, _input: &Stub) {}
346
347    fn visit_import(&mut self, input: &Program) {
348        self.visit_program(input)
349    }
350
351    fn visit_struct(&mut self, input: &Composite) {
352        input.const_parameters.iter().for_each(|input| self.visit_type(&input.type_));
353        input.members.iter().for_each(|member| self.visit_type(&member.type_));
354    }
355
356    fn visit_mapping(&mut self, input: &Mapping) {
357        self.visit_type(&input.key_type);
358        self.visit_type(&input.value_type);
359    }
360
361    fn visit_function(&mut self, input: &Function) {
362        input.const_parameters.iter().for_each(|input| self.visit_type(&input.type_));
363        input.input.iter().for_each(|input| self.visit_type(&input.type_));
364        input.output.iter().for_each(|output| self.visit_type(&output.type_));
365        self.visit_type(&input.output_type);
366        self.visit_block(&input.block);
367    }
368
369    fn visit_constructor(&mut self, input: &Constructor) {
370        self.visit_block(&input.block);
371    }
372
373    fn visit_function_stub(&mut self, _input: &FunctionStub) {}
374
375    fn visit_struct_stub(&mut self, _input: &Composite) {}
376}
leo_ast/passes/visitor.rs

leo_ast/passes/
visitor.rs
