air_parser/ast/
visit.rs

1//! This module provides infrastructure for the visitor pattern over the AirScript AST
2//!
3//! Implementations of [VisitMut] need only provide implementations of trait functions for
4//! the specific AST nodes which they are interested in, or which provide important context.
5//! By default, with no trait functions overridden, a visitor will simply traverse the AST
6//! top-down. When you override one of the trait functions, it is up to the implementation to
7//! drive the visitor down to children of the corresponding node type, if desired. For that purpose,
8//! this module exposes a number of `visit_mut_*` functions which can be called to perform the
9//! default visitor traversal for that node.
10use core::ops::ControlFlow;
11
12use crate::ast;
13
14/// This trait represents a mutable visitor over the AST.
15///
16/// See the example below for usage.
17///
18/// Each node visitor returns a `ControlFlow<T>`, which allows implementations to terminate the
19/// traversal early, which is particularly useful for error handling, but can be used for other
20/// purposes as well.
21///
22/// ## Example
23///
24/// ```rust
25/// use std::ops::ControlFlow;
26///
27/// use miden_diagnostics::{Span, Spanned};
28///
29/// use air_parser::ast::{self, visit};
30///
31/// /// A simple visitor which replaces accesses to constant values with the values themselves,
32/// /// evaluates constant expressions (i.e. expressions whose operands are constant), and propagates
33/// /// constants through let-bound variables (i.e. let bindings which are constant get replaced with
34/// /// the constant itself).
35/// struct ConstantPropagationVisitor {
36///     constants: std::collections::HashMap<ast::Identifier, Span<ast::ConstantExpr>>,
37/// }
38/// impl visit::VisitMut<()> for ConstantPropagationVisitor {
39///     // We override the visitor for constants so that we can record all of the known constant values
40///     fn visit_mut_constant(&mut self, constant: &mut ast::Constant) -> ControlFlow<()> {
41///         debug_assert_eq!(self.constants.get(&constant.name), None);
42///         let span = constant.span();
43///         self.constants.insert(constant.name, Span::new(span, constant.value.clone()));
44///         ControlFlow::Continue(())
45///     }
46///
47///     // We override the visitor for scalar expressions to propagate constants by evaluating any expressions
48///     // whose values or operands are all constant.
49///     fn visit_mut_scalar_expr(&mut self, expr: &mut ast::ScalarExpr) -> ControlFlow<()> {
50///         let span = expr.span();
51///         match expr {
52///             ast::ScalarExpr::Const(_) => ControlFlow::Continue(()),
53///             ast::ScalarExpr::SymbolAccess(sym) => {
54///                 let constant_value = self.constants.get(sym.name.as_ref()).cloned();
55///                 match constant_value.map(|s| (s.span(), s.item)){
56///                     None => (),
57///                     Some((span, ast::ConstantExpr::Scalar(value))) => {
58///                         assert_eq!(sym.access_type, ast::AccessType::Default);
59///                         core::mem::replace(expr, ast::ScalarExpr::Const(Span::new(span, value)));
60///                     }
61///                     Some((span, ast::ConstantExpr::Vector(value))) => {
62///                         match sym.access_type {
63///                             ast::AccessType::Index(idx) => {
64///                                 core::mem::replace(expr, ast::ScalarExpr::Const(Span::new(span, value[idx])));
65///                             }
66///                             _ => panic!("invalid constant reference, expected scalar access"),
67///                         }
68///                     }
69///                     Some((span, ast::ConstantExpr::Matrix(value))) => {
70///                         match sym.access_type {
71///                             ast::AccessType::Matrix(row, col) => {
72///                                 core::mem::replace(expr, ast::ScalarExpr::Const(Span::new(span, value[row][col])));
73///                             }
74///                             _ => panic!("invalid constant reference, expected scalar access"),
75///                         }
76///                     }
77///                 }
78///                 ControlFlow::Continue(())
79///             }
80///             ast::ScalarExpr::Binary(ast::BinaryExpr { op: ast::BinaryOp::Add, lhs, rhs, .. }) => {
81///                 visit::visit_mut_scalar_expr(self, lhs)?;
82///                 visit::visit_mut_scalar_expr(self, rhs)?;
83///                 // If both operands are constant, evaluate to a scalar constant
84///                 if let (ast::ScalarExpr::Const(l), ast::ScalarExpr::Const(r)) = (lhs.as_mut(), rhs.as_mut()) {
85///                     let folded = l.item + r.item;
86///                     core::mem::replace(expr, ast::ScalarExpr::Const(Span::new(span, folded)));
87///                 }
88///                 ControlFlow::Continue(())
89///             }
90///             /// The other arithmetic ops are basically the same as above
91///             _ => unimplemented!(),
92///         }
93///     }
94///
95///     // The implementation of this visitor is left as an exercise for the reader, but would be necessary
96///     // to ensure that we propagate constants through let-bound variables whose expressions are constant.
97///     //
98///     // It would also be necessary for us to preserve and update the bindings map in our visitor upon entry
99///     // to the let, and replace the original on exit from the let. This would automatically ensure that constants
100///     // get propagated through let-bound variables. All other identifiers can be safely ignored, with the exception
101///     // of those which shadow a binding that was previously constant, in which case it effectively erases that binding
102///     // from view.
103///     fn visit_mut_let(&mut self, _expr: &mut ast::Let) -> ControlFlow<()> {
104///         todo!()
105///     }
106/// }
107/// ```
108///
109pub trait VisitMut<T> {
110    fn visit_mut_module(&mut self, module: &mut ast::Module) -> ControlFlow<T> {
111        visit_mut_module(self, module)
112    }
113    fn visit_mut_import(&mut self, expr: &mut ast::Import) -> ControlFlow<T> {
114        visit_mut_import(self, expr)
115    }
116    fn visit_mut_constant(&mut self, expr: &mut ast::Constant) -> ControlFlow<T> {
117        visit_mut_constant(self, expr)
118    }
119    fn visit_mut_evaluator_function(
120        &mut self,
121        expr: &mut ast::EvaluatorFunction,
122    ) -> ControlFlow<T> {
123        visit_mut_evaluator_function(self, expr)
124    }
125    fn visit_mut_function(&mut self, expr: &mut ast::Function) -> ControlFlow<T> {
126        visit_mut_function(self, expr)
127    }
128    fn visit_mut_bus(&mut self, expr: &mut ast::Bus) -> ControlFlow<T> {
129        visit_mut_bus(self, expr)
130    }
131    fn visit_mut_periodic_column(&mut self, expr: &mut ast::PeriodicColumn) -> ControlFlow<T> {
132        visit_mut_periodic_column(self, expr)
133    }
134    fn visit_mut_public_input(&mut self, expr: &mut ast::PublicInput) -> ControlFlow<T> {
135        visit_mut_public_input(self, expr)
136    }
137    fn visit_mut_trace_segment(&mut self, expr: &mut ast::TraceSegment) -> ControlFlow<T> {
138        visit_mut_trace_segment(self, expr)
139    }
140    fn visit_mut_trace_binding(&mut self, expr: &mut ast::TraceBinding) -> ControlFlow<T> {
141        visit_mut_trace_binding(self, expr)
142    }
143    fn visit_mut_evaluator_trace_segment(
144        &mut self,
145        expr: &mut ast::TraceSegment,
146    ) -> ControlFlow<T> {
147        visit_mut_evaluator_trace_segment(self, expr)
148    }
149    fn visit_mut_evaluator_trace_binding(
150        &mut self,
151        expr: &mut ast::TraceBinding,
152    ) -> ControlFlow<T> {
153        visit_mut_evaluator_trace_binding(self, expr)
154    }
155    fn visit_mut_statement_block(&mut self, expr: &mut Vec<ast::Statement>) -> ControlFlow<T> {
156        visit_mut_statement_block(self, expr)
157    }
158    fn visit_mut_statement(&mut self, expr: &mut ast::Statement) -> ControlFlow<T> {
159        visit_mut_statement(self, expr)
160    }
161    fn visit_mut_let(&mut self, expr: &mut ast::Let) -> ControlFlow<T> {
162        visit_mut_let(self, expr)
163    }
164    fn visit_mut_boundary_constraints(
165        &mut self,
166        exprs: &mut Vec<ast::Statement>,
167    ) -> ControlFlow<T> {
168        self.visit_mut_statement_block(exprs)
169    }
170    fn visit_mut_enforce(&mut self, expr: &mut ast::ScalarExpr) -> ControlFlow<T> {
171        visit_mut_scalar_expr(self, expr)
172    }
173    fn visit_mut_enforce_if(
174        &mut self,
175        expr: &mut ast::ScalarExpr,
176        selector: &mut ast::ScalarExpr,
177    ) -> ControlFlow<T> {
178        self.visit_mut_enforce(expr)?;
179        self.visit_mut_scalar_expr(selector)
180    }
181    fn visit_mut_enforce_all(&mut self, expr: &mut ast::ListComprehension) -> ControlFlow<T> {
182        self.visit_mut_list_comprehension(expr)
183    }
184    fn visit_mut_bus_enforce(&mut self, expr: &mut ast::ListComprehension) -> ControlFlow<T> {
185        self.visit_mut_list_comprehension(expr)
186    }
187    fn visit_mut_integrity_constraints(
188        &mut self,
189        exprs: &mut Vec<ast::Statement>,
190    ) -> ControlFlow<T> {
191        self.visit_mut_statement_block(exprs)
192    }
193    fn visit_mut_expr(&mut self, expr: &mut ast::Expr) -> ControlFlow<T> {
194        visit_mut_expr(self, expr)
195    }
196    fn visit_mut_scalar_expr(&mut self, expr: &mut ast::ScalarExpr) -> ControlFlow<T> {
197        visit_mut_scalar_expr(self, expr)
198    }
199    fn visit_mut_binary_expr(&mut self, expr: &mut ast::BinaryExpr) -> ControlFlow<T> {
200        visit_mut_binary_expr(self, expr)
201    }
202    fn visit_mut_list_comprehension(
203        &mut self,
204        expr: &mut ast::ListComprehension,
205    ) -> ControlFlow<T> {
206        visit_mut_list_comprehension(self, expr)
207    }
208    fn visit_mut_call(&mut self, expr: &mut ast::Call) -> ControlFlow<T> {
209        visit_mut_call(self, expr)
210    }
211    fn visit_mut_bus_operation(&mut self, expr: &mut ast::BusOperation) -> ControlFlow<T> {
212        visit_mut_bus_operation(self, expr)
213    }
214    fn visit_mut_range_bound(&mut self, expr: &mut ast::RangeBound) -> ControlFlow<T> {
215        visit_mut_range_bound(self, expr)
216    }
217    fn visit_mut_access_type(&mut self, expr: &mut ast::AccessType) -> ControlFlow<T> {
218        visit_mut_access_type(self, expr)
219    }
220    fn visit_mut_const_symbol_access(
221        &mut self,
222        expr: &mut ast::ConstSymbolAccess,
223    ) -> ControlFlow<T> {
224        visit_mut_const_symbol_access(self, expr)
225    }
226    fn visit_mut_bounded_symbol_access(
227        &mut self,
228        expr: &mut ast::BoundedSymbolAccess,
229    ) -> ControlFlow<T> {
230        visit_mut_bounded_symbol_access(self, expr)
231    }
232    fn visit_mut_symbol_access(&mut self, expr: &mut ast::SymbolAccess) -> ControlFlow<T> {
233        visit_mut_symbol_access(self, expr)
234    }
235    fn visit_mut_resolvable_identifier(
236        &mut self,
237        expr: &mut ast::ResolvableIdentifier,
238    ) -> ControlFlow<T> {
239        visit_mut_resolvable_identifier(self, expr)
240    }
241    fn visit_mut_identifier(&mut self, expr: &mut ast::Identifier) -> ControlFlow<T> {
242        visit_mut_identifier(self, expr)
243    }
244    fn visit_mut_typed_identifier(
245        &mut self,
246        expr: &mut (ast::Identifier, ast::Type),
247    ) -> ControlFlow<T> {
248        visit_mut_typed_identifier(self, expr)
249    }
250}
251
252impl<V, T> VisitMut<T> for &mut V
253where
254    V: ?Sized + VisitMut<T>,
255{
256    fn visit_mut_module(&mut self, module: &mut ast::Module) -> ControlFlow<T> {
257        (**self).visit_mut_module(module)
258    }
259    fn visit_mut_import(&mut self, expr: &mut ast::Import) -> ControlFlow<T> {
260        (**self).visit_mut_import(expr)
261    }
262    fn visit_mut_constant(&mut self, expr: &mut ast::Constant) -> ControlFlow<T> {
263        (**self).visit_mut_constant(expr)
264    }
265    fn visit_mut_evaluator_function(
266        &mut self,
267        expr: &mut ast::EvaluatorFunction,
268    ) -> ControlFlow<T> {
269        (**self).visit_mut_evaluator_function(expr)
270    }
271    fn visit_mut_function(&mut self, expr: &mut ast::Function) -> ControlFlow<T> {
272        (**self).visit_mut_function(expr)
273    }
274    fn visit_mut_bus(&mut self, expr: &mut ast::Bus) -> ControlFlow<T> {
275        (**self).visit_mut_bus(expr)
276    }
277    fn visit_mut_periodic_column(&mut self, expr: &mut ast::PeriodicColumn) -> ControlFlow<T> {
278        (**self).visit_mut_periodic_column(expr)
279    }
280    fn visit_mut_public_input(&mut self, expr: &mut ast::PublicInput) -> ControlFlow<T> {
281        (**self).visit_mut_public_input(expr)
282    }
283    fn visit_mut_trace_segment(&mut self, expr: &mut ast::TraceSegment) -> ControlFlow<T> {
284        (**self).visit_mut_trace_segment(expr)
285    }
286    fn visit_mut_trace_binding(&mut self, expr: &mut ast::TraceBinding) -> ControlFlow<T> {
287        (**self).visit_mut_trace_binding(expr)
288    }
289    fn visit_mut_evaluator_trace_segment(
290        &mut self,
291        expr: &mut ast::TraceSegment,
292    ) -> ControlFlow<T> {
293        (**self).visit_mut_evaluator_trace_segment(expr)
294    }
295    fn visit_mut_evaluator_trace_binding(
296        &mut self,
297        expr: &mut ast::TraceBinding,
298    ) -> ControlFlow<T> {
299        (**self).visit_mut_evaluator_trace_binding(expr)
300    }
301    fn visit_mut_statement_block(&mut self, expr: &mut Vec<ast::Statement>) -> ControlFlow<T> {
302        (**self).visit_mut_statement_block(expr)
303    }
304    fn visit_mut_statement(&mut self, expr: &mut ast::Statement) -> ControlFlow<T> {
305        (**self).visit_mut_statement(expr)
306    }
307    fn visit_mut_let(&mut self, expr: &mut ast::Let) -> ControlFlow<T> {
308        (**self).visit_mut_let(expr)
309    }
310    fn visit_mut_boundary_constraints(
311        &mut self,
312        exprs: &mut Vec<ast::Statement>,
313    ) -> ControlFlow<T> {
314        (**self).visit_mut_boundary_constraints(exprs)
315    }
316    fn visit_mut_integrity_constraints(
317        &mut self,
318        exprs: &mut Vec<ast::Statement>,
319    ) -> ControlFlow<T> {
320        (**self).visit_mut_integrity_constraints(exprs)
321    }
322    fn visit_mut_enforce(&mut self, expr: &mut ast::ScalarExpr) -> ControlFlow<T> {
323        (**self).visit_mut_enforce(expr)
324    }
325    fn visit_mut_enforce_if(
326        &mut self,
327        expr: &mut ast::ScalarExpr,
328        selector: &mut ast::ScalarExpr,
329    ) -> ControlFlow<T> {
330        (**self).visit_mut_enforce_if(expr, selector)
331    }
332    fn visit_mut_enforce_all(&mut self, expr: &mut ast::ListComprehension) -> ControlFlow<T> {
333        (**self).visit_mut_enforce_all(expr)
334    }
335    fn visit_mut_bus_enforce(&mut self, expr: &mut ast::ListComprehension) -> ControlFlow<T> {
336        (**self).visit_mut_bus_enforce(expr)
337    }
338    fn visit_mut_expr(&mut self, expr: &mut ast::Expr) -> ControlFlow<T> {
339        (**self).visit_mut_expr(expr)
340    }
341    fn visit_mut_scalar_expr(&mut self, expr: &mut ast::ScalarExpr) -> ControlFlow<T> {
342        (**self).visit_mut_scalar_expr(expr)
343    }
344    fn visit_mut_binary_expr(&mut self, expr: &mut ast::BinaryExpr) -> ControlFlow<T> {
345        (**self).visit_mut_binary_expr(expr)
346    }
347    fn visit_mut_list_comprehension(
348        &mut self,
349        expr: &mut ast::ListComprehension,
350    ) -> ControlFlow<T> {
351        (**self).visit_mut_list_comprehension(expr)
352    }
353    fn visit_mut_call(&mut self, expr: &mut ast::Call) -> ControlFlow<T> {
354        (**self).visit_mut_call(expr)
355    }
356    fn visit_mut_bus_operation(&mut self, expr: &mut ast::BusOperation) -> ControlFlow<T> {
357        (**self).visit_mut_bus_operation(expr)
358    }
359    fn visit_mut_range_bound(&mut self, expr: &mut ast::RangeBound) -> ControlFlow<T> {
360        (**self).visit_mut_range_bound(expr)
361    }
362    fn visit_mut_access_type(&mut self, expr: &mut ast::AccessType) -> ControlFlow<T> {
363        (**self).visit_mut_access_type(expr)
364    }
365    fn visit_mut_const_symbol_access(
366        &mut self,
367        expr: &mut ast::ConstSymbolAccess,
368    ) -> ControlFlow<T> {
369        (**self).visit_mut_const_symbol_access(expr)
370    }
371    fn visit_mut_bounded_symbol_access(
372        &mut self,
373        expr: &mut ast::BoundedSymbolAccess,
374    ) -> ControlFlow<T> {
375        (**self).visit_mut_bounded_symbol_access(expr)
376    }
377    fn visit_mut_symbol_access(&mut self, expr: &mut ast::SymbolAccess) -> ControlFlow<T> {
378        (**self).visit_mut_symbol_access(expr)
379    }
380    fn visit_mut_resolvable_identifier(
381        &mut self,
382        expr: &mut ast::ResolvableIdentifier,
383    ) -> ControlFlow<T> {
384        (**self).visit_mut_resolvable_identifier(expr)
385    }
386    fn visit_mut_identifier(&mut self, expr: &mut ast::Identifier) -> ControlFlow<T> {
387        (**self).visit_mut_identifier(expr)
388    }
389    fn visit_mut_typed_identifier(
390        &mut self,
391        expr: &mut (ast::Identifier, ast::Type),
392    ) -> ControlFlow<T> {
393        (**self).visit_mut_typed_identifier(expr)
394    }
395}
396
397pub fn visit_mut_module<V, T>(visitor: &mut V, module: &mut ast::Module) -> ControlFlow<T>
398where
399    V: ?Sized + VisitMut<T>,
400{
401    for import in module.imports.values_mut() {
402        visitor.visit_mut_import(import)?;
403    }
404    for constant in module.constants.values_mut() {
405        visitor.visit_mut_constant(constant)?;
406    }
407    for evaluator in module.evaluators.values_mut() {
408        visitor.visit_mut_evaluator_function(evaluator)?;
409    }
410    for function in module.functions.values_mut() {
411        visitor.visit_mut_function(function)?;
412    }
413    for bus in module.buses.values_mut() {
414        visitor.visit_mut_bus(bus)?;
415    }
416    for column in module.periodic_columns.values_mut() {
417        visitor.visit_mut_periodic_column(column)?;
418    }
419    for input in module.public_inputs.values_mut() {
420        visitor.visit_mut_public_input(input)?;
421    }
422    for segment in module.trace_columns.iter_mut() {
423        visitor.visit_mut_trace_segment(segment)?;
424    }
425    if let Some(bc) = module.boundary_constraints.as_mut() {
426        if !bc.is_empty() {
427            visitor.visit_mut_boundary_constraints(bc)?;
428        }
429    }
430    if let Some(ic) = module.integrity_constraints.as_mut() {
431        if !ic.is_empty() {
432            visitor.visit_mut_integrity_constraints(ic)?;
433        }
434    }
435
436    ControlFlow::Continue(())
437}
438
439pub fn visit_mut_import<V, T>(_visitor: &mut V, _expr: &mut ast::Import) -> ControlFlow<T>
440where
441    V: ?Sized + VisitMut<T>,
442{
443    ControlFlow::Continue(())
444}
445
446pub fn visit_mut_constant<V, T>(visitor: &mut V, expr: &mut ast::Constant) -> ControlFlow<T>
447where
448    V: ?Sized + VisitMut<T>,
449{
450    visitor.visit_mut_identifier(&mut expr.name)
451}
452
453pub fn visit_mut_trace_segment<V, T>(
454    visitor: &mut V,
455    expr: &mut ast::TraceSegment,
456) -> ControlFlow<T>
457where
458    V: ?Sized + VisitMut<T>,
459{
460    for binding in expr.bindings.iter_mut() {
461        visitor.visit_mut_trace_binding(binding)?;
462    }
463    ControlFlow::Continue(())
464}
465
466pub fn visit_mut_trace_binding<V, T>(
467    visitor: &mut V,
468    expr: &mut ast::TraceBinding,
469) -> ControlFlow<T>
470where
471    V: ?Sized + VisitMut<T>,
472{
473    if let Some(name) = expr.name.as_mut() {
474        visitor.visit_mut_identifier(name)?;
475    }
476    ControlFlow::Continue(())
477}
478
479pub fn visit_mut_evaluator_function<V, T>(
480    visitor: &mut V,
481    expr: &mut ast::EvaluatorFunction,
482) -> ControlFlow<T>
483where
484    V: ?Sized + VisitMut<T>,
485{
486    visitor.visit_mut_identifier(&mut expr.name)?;
487    for segment in expr.params.iter_mut() {
488        visitor.visit_mut_evaluator_trace_segment(segment)?;
489    }
490    visitor.visit_mut_statement_block(&mut expr.body)
491}
492
493pub fn visit_mut_function<V, T>(visitor: &mut V, expr: &mut ast::Function) -> ControlFlow<T>
494where
495    V: ?Sized + VisitMut<T>,
496{
497    visitor.visit_mut_identifier(&mut expr.name)?;
498    for param in expr.params.iter_mut() {
499        visitor.visit_mut_typed_identifier(param)?;
500    }
501    visitor.visit_mut_statement_block(&mut expr.body)
502}
503
504pub fn visit_mut_bus<V, T>(visitor: &mut V, expr: &mut ast::Bus) -> ControlFlow<T>
505where
506    V: ?Sized + VisitMut<T>,
507{
508    visitor.visit_mut_identifier(&mut expr.name)
509}
510
511pub fn visit_mut_evaluator_trace_segment<V, T>(
512    visitor: &mut V,
513    expr: &mut ast::TraceSegment,
514) -> ControlFlow<T>
515where
516    V: ?Sized + VisitMut<T>,
517{
518    for binding in expr.bindings.iter_mut() {
519        visitor.visit_mut_evaluator_trace_binding(binding)?;
520    }
521    ControlFlow::Continue(())
522}
523
524pub fn visit_mut_evaluator_trace_binding<V, T>(
525    visitor: &mut V,
526    expr: &mut ast::TraceBinding,
527) -> ControlFlow<T>
528where
529    V: ?Sized + VisitMut<T>,
530{
531    if let Some(name) = expr.name.as_mut() {
532        visitor.visit_mut_identifier(name)?;
533    }
534    ControlFlow::Continue(())
535}
536
537pub fn visit_mut_periodic_column<V, T>(
538    visitor: &mut V,
539    expr: &mut ast::PeriodicColumn,
540) -> ControlFlow<T>
541where
542    V: ?Sized + VisitMut<T>,
543{
544    visitor.visit_mut_identifier(&mut expr.name)
545}
546
547pub fn visit_mut_public_input<V, T>(visitor: &mut V, expr: &mut ast::PublicInput) -> ControlFlow<T>
548where
549    V: ?Sized + VisitMut<T>,
550{
551    visitor.visit_mut_identifier(&mut expr.name())
552}
553
554pub fn visit_mut_statement_block<V, T>(
555    visitor: &mut V,
556    statements: &mut [ast::Statement],
557) -> ControlFlow<T>
558where
559    V: ?Sized + VisitMut<T>,
560{
561    for statement in statements.iter_mut() {
562        visitor.visit_mut_statement(statement)?;
563    }
564    ControlFlow::Continue(())
565}
566
567pub fn visit_mut_statement<V, T>(visitor: &mut V, expr: &mut ast::Statement) -> ControlFlow<T>
568where
569    V: ?Sized + VisitMut<T>,
570{
571    match expr {
572        ast::Statement::Let(expr) => visitor.visit_mut_let(expr),
573        ast::Statement::Enforce(expr) => visitor.visit_mut_enforce(expr),
574        ast::Statement::EnforceIf(expr, selector) => visitor.visit_mut_enforce_if(expr, selector),
575        ast::Statement::EnforceAll(expr) => visitor.visit_mut_enforce_all(expr),
576        ast::Statement::Expr(expr) => visitor.visit_mut_expr(expr),
577        ast::Statement::BusEnforce(expr) => visitor.visit_mut_bus_enforce(expr),
578    }
579}
580
581pub fn visit_mut_let<V, T>(visitor: &mut V, expr: &mut ast::Let) -> ControlFlow<T>
582where
583    V: ?Sized + VisitMut<T>,
584{
585    visitor.visit_mut_expr(&mut expr.value)?;
586    visitor.visit_mut_identifier(&mut expr.name)?;
587    for statement in expr.body.iter_mut() {
588        visitor.visit_mut_statement(statement)?;
589    }
590    ControlFlow::Continue(())
591}
592
593pub fn visit_mut_expr<V, T>(visitor: &mut V, expr: &mut ast::Expr) -> ControlFlow<T>
594where
595    V: ?Sized + VisitMut<T>,
596{
597    match expr {
598        ast::Expr::Const(_) => ControlFlow::Continue(()),
599        ast::Expr::Range(range) => {
600            visitor.visit_mut_range_bound(&mut range.start)?;
601            visitor.visit_mut_range_bound(&mut range.end)?;
602            ControlFlow::Continue(())
603        }
604        ast::Expr::Vector(exprs) => {
605            for expr in exprs.iter_mut() {
606                visitor.visit_mut_expr(expr)?;
607            }
608            ControlFlow::Continue(())
609        }
610        ast::Expr::Matrix(matrix) => {
611            for exprs in matrix.iter_mut() {
612                for expr in exprs.iter_mut() {
613                    visitor.visit_mut_scalar_expr(expr)?;
614                }
615            }
616            ControlFlow::Continue(())
617        }
618        ast::Expr::SymbolAccess(expr) => visitor.visit_mut_symbol_access(expr),
619        ast::Expr::Binary(expr) => visitor.visit_mut_binary_expr(expr),
620        ast::Expr::Call(expr) => visitor.visit_mut_call(expr),
621        ast::Expr::ListComprehension(expr) => visitor.visit_mut_list_comprehension(expr),
622        ast::Expr::Let(expr) => visitor.visit_mut_let(expr),
623        ast::Expr::BusOperation(expr) => visitor.visit_mut_bus_operation(expr),
624        ast::Expr::Null(_) | ast::Expr::Unconstrained(_) => ControlFlow::Continue(()),
625    }
626}
627
628pub fn visit_mut_scalar_expr<V, T>(visitor: &mut V, expr: &mut ast::ScalarExpr) -> ControlFlow<T>
629where
630    V: ?Sized + VisitMut<T>,
631{
632    match expr {
633        ast::ScalarExpr::Const(_)
634        | ast::ScalarExpr::Null(_)
635        | ast::ScalarExpr::Unconstrained(_) => ControlFlow::Continue(()),
636        ast::ScalarExpr::SymbolAccess(expr) => visitor.visit_mut_symbol_access(expr),
637        ast::ScalarExpr::BoundedSymbolAccess(expr) => visitor.visit_mut_bounded_symbol_access(expr),
638        ast::ScalarExpr::Binary(expr) => visitor.visit_mut_binary_expr(expr),
639        ast::ScalarExpr::Call(expr) => visitor.visit_mut_call(expr),
640        ast::ScalarExpr::Let(expr) => visitor.visit_mut_let(expr),
641        ast::ScalarExpr::BusOperation(expr) => visitor.visit_mut_bus_operation(expr),
642    }
643}
644
645pub fn visit_mut_binary_expr<V, T>(visitor: &mut V, expr: &mut ast::BinaryExpr) -> ControlFlow<T>
646where
647    V: ?Sized + VisitMut<T>,
648{
649    visitor.visit_mut_scalar_expr(expr.lhs.as_mut())?;
650    visitor.visit_mut_scalar_expr(expr.rhs.as_mut())
651}
652
653pub fn visit_mut_list_comprehension<V, T>(
654    visitor: &mut V,
655    expr: &mut ast::ListComprehension,
656) -> ControlFlow<T>
657where
658    V: ?Sized + VisitMut<T>,
659{
660    for binding in expr.bindings.iter_mut() {
661        visitor.visit_mut_identifier(binding)?;
662    }
663    for iterable in expr.iterables.iter_mut() {
664        visitor.visit_mut_expr(iterable)?;
665    }
666    if let Some(selector) = expr.selector.as_mut() {
667        visitor.visit_mut_scalar_expr(selector)?;
668    }
669    visitor.visit_mut_scalar_expr(expr.body.as_mut())
670}
671
672pub fn visit_mut_call<V, T>(visitor: &mut V, expr: &mut ast::Call) -> ControlFlow<T>
673where
674    V: ?Sized + VisitMut<T>,
675{
676    visitor.visit_mut_resolvable_identifier(&mut expr.callee)?;
677    for arg in expr.args.iter_mut() {
678        visitor.visit_mut_expr(arg)?;
679    }
680    ControlFlow::Continue(())
681}
682
683pub fn visit_mut_bus_operation<V, T>(
684    visitor: &mut V,
685    expr: &mut ast::BusOperation,
686) -> ControlFlow<T>
687where
688    V: ?Sized + VisitMut<T>,
689{
690    visitor.visit_mut_resolvable_identifier(&mut expr.bus)?;
691    for arg in expr.args.iter_mut() {
692        visitor.visit_mut_expr(arg)?;
693    }
694    ControlFlow::Continue(())
695}
696
697pub fn visit_mut_range_bound<V, T>(visitor: &mut V, expr: &mut ast::RangeBound) -> ControlFlow<T>
698where
699    V: ?Sized + VisitMut<T>,
700{
701    match expr {
702        ast::RangeBound::Const(_) => ControlFlow::Continue(()),
703        ast::RangeBound::SymbolAccess(access) => visitor.visit_mut_const_symbol_access(access),
704    }
705}
706
707pub fn visit_mut_access_type<V, T>(visitor: &mut V, expr: &mut ast::AccessType) -> ControlFlow<T>
708where
709    V: ?Sized + VisitMut<T>,
710{
711    match expr {
712        ast::AccessType::Default | ast::AccessType::Index(_) | ast::AccessType::Matrix(_, _) => {
713            ControlFlow::Continue(())
714        }
715        ast::AccessType::Slice(range) => {
716            visitor.visit_mut_range_bound(&mut range.start)?;
717            visitor.visit_mut_range_bound(&mut range.end)
718        }
719    }
720}
721
722pub fn visit_mut_const_symbol_access<V, T>(
723    visitor: &mut V,
724    expr: &mut ast::ConstSymbolAccess,
725) -> ControlFlow<T>
726where
727    V: ?Sized + VisitMut<T>,
728{
729    visitor.visit_mut_resolvable_identifier(&mut expr.name)
730}
731
732pub fn visit_mut_bounded_symbol_access<V, T>(
733    visitor: &mut V,
734    expr: &mut ast::BoundedSymbolAccess,
735) -> ControlFlow<T>
736where
737    V: ?Sized + VisitMut<T>,
738{
739    visitor.visit_mut_symbol_access(&mut expr.column)
740}
741
742pub fn visit_mut_symbol_access<V, T>(
743    visitor: &mut V,
744    expr: &mut ast::SymbolAccess,
745) -> ControlFlow<T>
746where
747    V: ?Sized + VisitMut<T>,
748{
749    visitor.visit_mut_resolvable_identifier(&mut expr.name)
750}
751
752pub fn visit_mut_resolvable_identifier<V, T>(
753    _visitor: &mut V,
754    _expr: &mut ast::ResolvableIdentifier,
755) -> ControlFlow<T>
756where
757    V: ?Sized + VisitMut<T>,
758{
759    ControlFlow::Continue(())
760}
761
762pub fn visit_mut_identifier<V, T>(_visitor: &mut V, _expr: &mut ast::Identifier) -> ControlFlow<T>
763where
764    V: ?Sized + VisitMut<T>,
765{
766    ControlFlow::Continue(())
767}
768
769pub fn visit_mut_typed_identifier<V, T>(
770    _visitor: &mut V,
771    _expr: &mut (ast::Identifier, ast::Type),
772) -> ControlFlow<T>
773where
774    V: ?Sized + VisitMut<T>,
775{
776    ControlFlow::Continue(())
777}
air_parser/ast/visit.rs

air_parser/ast/
visit.rs
