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 expressions in the AST.
46pub trait ExpressionVisitor {
47    type AdditionalInput: Default;
48    type Output: Default;
49
50    fn visit_expression(&mut self, input: &Expression, additional: &Self::AdditionalInput) -> Self::Output {
51        match input {
52            Expression::Array(array) => self.visit_array(array, additional),
53            Expression::ArrayAccess(access) => self.visit_array_access(access, additional),
54            Expression::AssociatedConstant(constant) => self.visit_associated_constant(constant, additional),
55            Expression::AssociatedFunction(function) => self.visit_associated_function(function, additional),
56            Expression::Binary(binary) => self.visit_binary(binary, additional),
57            Expression::Call(call) => self.visit_call(call, additional),
58            Expression::Cast(cast) => self.visit_cast(cast, additional),
59            Expression::Struct(struct_) => self.visit_struct_init(struct_, additional),
60            Expression::Err(err) => self.visit_err(err, additional),
61            Expression::Identifier(identifier) => self.visit_identifier(identifier, additional),
62            Expression::Literal(literal) => self.visit_literal(literal, additional),
63            Expression::Locator(locator) => self.visit_locator(locator, additional),
64            Expression::MemberAccess(access) => self.visit_member_access(access, additional),
65            Expression::Ternary(ternary) => self.visit_ternary(ternary, additional),
66            Expression::Tuple(tuple) => self.visit_tuple(tuple, additional),
67            Expression::TupleAccess(access) => self.visit_tuple_access(access, additional),
68            Expression::Unary(unary) => self.visit_unary(unary, additional),
69            Expression::Unit(unit) => self.visit_unit(unit, additional),
70        }
71    }
72
73    fn visit_array_access(&mut self, input: &ArrayAccess, additional: &Self::AdditionalInput) -> Self::Output {
74        self.visit_expression(&input.array, additional);
75        self.visit_expression(&input.index, additional);
76        Default::default()
77    }
78
79    fn visit_member_access(&mut self, input: &MemberAccess, additional: &Self::AdditionalInput) -> Self::Output {
80        self.visit_expression(&input.inner, additional);
81        Default::default()
82    }
83
84    fn visit_tuple_access(&mut self, input: &TupleAccess, additional: &Self::AdditionalInput) -> Self::Output {
85        self.visit_expression(&input.tuple, additional);
86        Default::default()
87    }
88
89    fn visit_array(&mut self, input: &ArrayExpression, additional: &Self::AdditionalInput) -> Self::Output {
90        input.elements.iter().for_each(|expr| {
91            self.visit_expression(expr, additional);
92        });
93        Default::default()
94    }
95
96    fn visit_associated_constant(
97        &mut self,
98        _input: &AssociatedConstantExpression,
99        _additional: &Self::AdditionalInput,
100    ) -> Self::Output {
101        Default::default()
102    }
103
104    fn visit_associated_function(
105        &mut self,
106        input: &AssociatedFunctionExpression,
107        _additional: &Self::AdditionalInput,
108    ) -> Self::Output {
109        input.arguments.iter().for_each(|arg| {
110            self.visit_expression(arg, &Default::default());
111        });
112        Default::default()
113    }
114
115    fn visit_binary(&mut self, input: &BinaryExpression, additional: &Self::AdditionalInput) -> Self::Output {
116        self.visit_expression(&input.left, additional);
117        self.visit_expression(&input.right, additional);
118        Default::default()
119    }
120
121    fn visit_call(&mut self, input: &CallExpression, additional: &Self::AdditionalInput) -> Self::Output {
122        input.arguments.iter().for_each(|expr| {
123            self.visit_expression(expr, additional);
124        });
125        Default::default()
126    }
127
128    fn visit_cast(&mut self, input: &CastExpression, additional: &Self::AdditionalInput) -> Self::Output {
129        self.visit_expression(&input.expression, additional);
130        Default::default()
131    }
132
133    fn visit_struct_init(&mut self, input: &StructExpression, additional: &Self::AdditionalInput) -> Self::Output {
134        for StructVariableInitializer { expression, .. } in input.members.iter() {
135            if let Some(expression) = expression {
136                self.visit_expression(expression, additional);
137            }
138        }
139        Default::default()
140    }
141
142    fn visit_err(&mut self, _input: &ErrExpression, _additional: &Self::AdditionalInput) -> Self::Output {
143        panic!("`ErrExpression`s should not be in the AST at this phase of compilation.")
144    }
145
146    fn visit_identifier(&mut self, _input: &Identifier, _additional: &Self::AdditionalInput) -> Self::Output {
147        Default::default()
148    }
149
150    fn visit_literal(&mut self, _input: &Literal, _additional: &Self::AdditionalInput) -> Self::Output {
151        Default::default()
152    }
153
154    fn visit_locator(&mut self, _input: &LocatorExpression, _additional: &Self::AdditionalInput) -> Self::Output {
155        Default::default()
156    }
157
158    fn visit_ternary(&mut self, input: &TernaryExpression, additional: &Self::AdditionalInput) -> Self::Output {
159        self.visit_expression(&input.condition, additional);
160        self.visit_expression(&input.if_true, additional);
161        self.visit_expression(&input.if_false, additional);
162        Default::default()
163    }
164
165    fn visit_tuple(&mut self, input: &TupleExpression, additional: &Self::AdditionalInput) -> Self::Output {
166        input.elements.iter().for_each(|expr| {
167            self.visit_expression(expr, additional);
168        });
169        Default::default()
170    }
171
172    fn visit_unary(&mut self, input: &UnaryExpression, additional: &Self::AdditionalInput) -> Self::Output {
173        self.visit_expression(&input.receiver, additional);
174        Default::default()
175    }
176
177    fn visit_unit(&mut self, _input: &UnitExpression, _additional: &Self::AdditionalInput) -> Self::Output {
178        Default::default()
179    }
180}
181
182/// A Visitor trait for statements in the AST.
183pub trait StatementVisitor: ExpressionVisitor {
184    fn visit_statement(&mut self, input: &Statement) {
185        match input {
186            Statement::Assert(stmt) => self.visit_assert(stmt),
187            Statement::Assign(stmt) => self.visit_assign(stmt),
188            Statement::Block(stmt) => self.visit_block(stmt),
189            Statement::Conditional(stmt) => self.visit_conditional(stmt),
190            Statement::Const(stmt) => self.visit_const(stmt),
191            Statement::Definition(stmt) => self.visit_definition(stmt),
192            Statement::Expression(stmt) => self.visit_expression_statement(stmt),
193            Statement::Iteration(stmt) => self.visit_iteration(stmt),
194            Statement::Return(stmt) => self.visit_return(stmt),
195        }
196    }
197
198    fn visit_assert(&mut self, input: &AssertStatement) {
199        match &input.variant {
200            AssertVariant::Assert(expr) => self.visit_expression(expr, &Default::default()),
201            AssertVariant::AssertEq(left, right) | AssertVariant::AssertNeq(left, right) => {
202                self.visit_expression(left, &Default::default());
203                self.visit_expression(right, &Default::default())
204            }
205        };
206    }
207
208    fn visit_assign(&mut self, input: &AssignStatement) {
209        self.visit_expression(&input.value, &Default::default());
210    }
211
212    fn visit_block(&mut self, input: &Block) {
213        input.statements.iter().for_each(|stmt| self.visit_statement(stmt));
214    }
215
216    fn visit_conditional(&mut self, input: &ConditionalStatement) {
217        self.visit_expression(&input.condition, &Default::default());
218        self.visit_block(&input.then);
219        if let Some(stmt) = input.otherwise.as_ref() {
220            self.visit_statement(stmt);
221        }
222    }
223
224    fn visit_const(&mut self, input: &ConstDeclaration) {
225        self.visit_expression(&input.value, &Default::default());
226    }
227
228    fn visit_definition(&mut self, input: &DefinitionStatement) {
229        self.visit_expression(&input.value, &Default::default());
230    }
231
232    fn visit_expression_statement(&mut self, input: &ExpressionStatement) {
233        self.visit_expression(&input.expression, &Default::default());
234    }
235
236    fn visit_iteration(&mut self, input: &IterationStatement) {
237        self.visit_expression(&input.start, &Default::default());
238        self.visit_expression(&input.stop, &Default::default());
239        self.visit_block(&input.block);
240    }
241
242    fn visit_return(&mut self, input: &ReturnStatement) {
243        self.visit_expression(&input.expression, &Default::default());
244    }
245}
246
247/// A Visitor trait for the program represented by the AST.
248pub trait ProgramVisitor: StatementVisitor {
249    fn visit_program(&mut self, input: &Program) {
250        input.imports.values().for_each(|import| self.visit_import(&import.0));
251        input.stubs.values().for_each(|stub| self.visit_stub(stub));
252        input.program_scopes.values().for_each(|scope| self.visit_program_scope(scope));
253    }
254
255    fn visit_program_scope(&mut self, input: &ProgramScope) {
256        input.structs.iter().for_each(|(_, c)| (self.visit_struct(c)));
257
258        input.mappings.iter().for_each(|(_, c)| (self.visit_mapping(c)));
259
260        input.functions.iter().for_each(|(_, c)| (self.visit_function(c)));
261
262        input.consts.iter().for_each(|(_, c)| (self.visit_const(c)));
263    }
264
265    fn visit_stub(&mut self, _input: &Stub) {}
266
267    fn visit_import(&mut self, input: &Program) {
268        self.visit_program(input)
269    }
270
271    fn visit_struct(&mut self, _input: &Composite) {}
272
273    fn visit_mapping(&mut self, _input: &Mapping) {}
274
275    fn visit_function(&mut self, input: &Function) {
276        self.visit_block(&input.block);
277    }
278
279    fn visit_function_stub(&mut self, _input: &FunctionStub) {}
280
281    fn visit_struct_stub(&mut self, _input: &Composite) {}
282}