forge_fmt/
visit.rs

1//! Visitor helpers to traverse the [solang Solidity Parse Tree](solang_parser::pt).
2
3use crate::solang_ext::{pt::*, CodeLocationExt};
4
5/// A trait that is invoked while traversing the Solidity Parse Tree.
6/// Each method of the [Visitor] trait is a hook that can be potentially overridden.
7///
8/// Currently the main implementer of this trait is the [`Formatter`](crate::Formatter<'_>) struct.
9pub trait Visitor {
10    type Error: std::error::Error;
11
12    fn visit_source(&mut self, _loc: Loc) -> Result<(), Self::Error> {
13        Ok(())
14    }
15
16    fn visit_source_unit(&mut self, _source_unit: &mut SourceUnit) -> Result<(), Self::Error> {
17        Ok(())
18    }
19
20    fn visit_contract(&mut self, _contract: &mut ContractDefinition) -> Result<(), Self::Error> {
21        Ok(())
22    }
23
24    fn visit_annotation(&mut self, annotation: &mut Annotation) -> Result<(), Self::Error> {
25        self.visit_source(annotation.loc)
26    }
27
28    fn visit_pragma(&mut self, pragma: &mut PragmaDirective) -> Result<(), Self::Error> {
29        self.visit_source(pragma.loc())
30    }
31
32    fn visit_import_plain(
33        &mut self,
34        _loc: Loc,
35        _import: &mut ImportPath,
36    ) -> Result<(), Self::Error> {
37        Ok(())
38    }
39
40    fn visit_import_global(
41        &mut self,
42        _loc: Loc,
43        _global: &mut ImportPath,
44        _alias: &mut Identifier,
45    ) -> Result<(), Self::Error> {
46        Ok(())
47    }
48
49    fn visit_import_renames(
50        &mut self,
51        _loc: Loc,
52        _imports: &mut [(Identifier, Option<Identifier>)],
53        _from: &mut ImportPath,
54    ) -> Result<(), Self::Error> {
55        Ok(())
56    }
57
58    fn visit_enum(&mut self, _enum: &mut EnumDefinition) -> Result<(), Self::Error> {
59        Ok(())
60    }
61
62    fn visit_assembly(
63        &mut self,
64        loc: Loc,
65        _dialect: &mut Option<StringLiteral>,
66        _block: &mut YulBlock,
67        _flags: &mut Option<Vec<StringLiteral>>,
68    ) -> Result<(), Self::Error> {
69        self.visit_source(loc)
70    }
71
72    fn visit_block(
73        &mut self,
74        loc: Loc,
75        _unchecked: bool,
76        _statements: &mut Vec<Statement>,
77    ) -> Result<(), Self::Error> {
78        self.visit_source(loc)
79    }
80
81    fn visit_args(&mut self, loc: Loc, _args: &mut Vec<NamedArgument>) -> Result<(), Self::Error> {
82        self.visit_source(loc)
83    }
84
85    /// Don't write semicolon at the end because expressions can appear as both
86    /// part of other node and a statement in the function body
87    fn visit_expr(&mut self, loc: Loc, _expr: &mut Expression) -> Result<(), Self::Error> {
88        self.visit_source(loc)
89    }
90
91    fn visit_ident(&mut self, loc: Loc, _ident: &mut Identifier) -> Result<(), Self::Error> {
92        self.visit_source(loc)
93    }
94
95    fn visit_ident_path(&mut self, idents: &mut IdentifierPath) -> Result<(), Self::Error> {
96        self.visit_source(idents.loc)
97    }
98
99    fn visit_emit(&mut self, loc: Loc, _event: &mut Expression) -> Result<(), Self::Error> {
100        self.visit_source(loc)?;
101        self.visit_stray_semicolon()?;
102
103        Ok(())
104    }
105
106    fn visit_var_definition(&mut self, var: &mut VariableDefinition) -> Result<(), Self::Error> {
107        self.visit_source(var.loc)?;
108        self.visit_stray_semicolon()?;
109
110        Ok(())
111    }
112
113    fn visit_var_definition_stmt(
114        &mut self,
115        loc: Loc,
116        _declaration: &mut VariableDeclaration,
117        _expr: &mut Option<Expression>,
118    ) -> Result<(), Self::Error> {
119        self.visit_source(loc)?;
120        self.visit_stray_semicolon()
121    }
122
123    fn visit_var_declaration(&mut self, var: &mut VariableDeclaration) -> Result<(), Self::Error> {
124        self.visit_source(var.loc)
125    }
126
127    fn visit_return(
128        &mut self,
129        loc: Loc,
130        _expr: &mut Option<Expression>,
131    ) -> Result<(), Self::Error> {
132        self.visit_source(loc)?;
133        self.visit_stray_semicolon()?;
134
135        Ok(())
136    }
137
138    fn visit_revert(
139        &mut self,
140        loc: Loc,
141        _error: &mut Option<IdentifierPath>,
142        _args: &mut Vec<Expression>,
143    ) -> Result<(), Self::Error> {
144        self.visit_source(loc)?;
145        self.visit_stray_semicolon()?;
146
147        Ok(())
148    }
149
150    fn visit_revert_named_args(
151        &mut self,
152        loc: Loc,
153        _error: &mut Option<IdentifierPath>,
154        _args: &mut Vec<NamedArgument>,
155    ) -> Result<(), Self::Error> {
156        self.visit_source(loc)?;
157        self.visit_stray_semicolon()?;
158
159        Ok(())
160    }
161
162    fn visit_break(&mut self, loc: Loc, _semicolon: bool) -> Result<(), Self::Error> {
163        self.visit_source(loc)
164    }
165
166    fn visit_continue(&mut self, loc: Loc, _semicolon: bool) -> Result<(), Self::Error> {
167        self.visit_source(loc)
168    }
169
170    #[expect(clippy::type_complexity)]
171    fn visit_try(
172        &mut self,
173        loc: Loc,
174        _expr: &mut Expression,
175        _returns: &mut Option<(Vec<(Loc, Option<Parameter>)>, Box<Statement>)>,
176        _clauses: &mut Vec<CatchClause>,
177    ) -> Result<(), Self::Error> {
178        self.visit_source(loc)
179    }
180
181    fn visit_if(
182        &mut self,
183        loc: Loc,
184        _cond: &mut Expression,
185        _if_branch: &mut Box<Statement>,
186        _else_branch: &mut Option<Box<Statement>>,
187        _is_first_stmt: bool,
188    ) -> Result<(), Self::Error> {
189        self.visit_source(loc)
190    }
191
192    fn visit_do_while(
193        &mut self,
194        loc: Loc,
195        _body: &mut Statement,
196        _cond: &mut Expression,
197    ) -> Result<(), Self::Error> {
198        self.visit_source(loc)
199    }
200
201    fn visit_while(
202        &mut self,
203        loc: Loc,
204        _cond: &mut Expression,
205        _body: &mut Statement,
206    ) -> Result<(), Self::Error> {
207        self.visit_source(loc)
208    }
209
210    fn visit_for(
211        &mut self,
212        loc: Loc,
213        _init: &mut Option<Box<Statement>>,
214        _cond: &mut Option<Box<Expression>>,
215        _update: &mut Option<Box<Expression>>,
216        _body: &mut Option<Box<Statement>>,
217    ) -> Result<(), Self::Error> {
218        self.visit_source(loc)
219    }
220
221    fn visit_function(&mut self, func: &mut FunctionDefinition) -> Result<(), Self::Error> {
222        self.visit_source(func.loc())?;
223        if func.body.is_none() {
224            self.visit_stray_semicolon()?;
225        }
226
227        Ok(())
228    }
229
230    fn visit_function_attribute(
231        &mut self,
232        attribute: &mut FunctionAttribute,
233    ) -> Result<(), Self::Error> {
234        self.visit_source(attribute.loc())?;
235        Ok(())
236    }
237
238    fn visit_var_attribute(
239        &mut self,
240        attribute: &mut VariableAttribute,
241    ) -> Result<(), Self::Error> {
242        self.visit_source(attribute.loc())?;
243        Ok(())
244    }
245
246    fn visit_base(&mut self, base: &mut Base) -> Result<(), Self::Error> {
247        self.visit_source(base.loc)
248    }
249
250    fn visit_parameter(&mut self, parameter: &mut Parameter) -> Result<(), Self::Error> {
251        self.visit_source(parameter.loc)
252    }
253
254    fn visit_struct(&mut self, structure: &mut StructDefinition) -> Result<(), Self::Error> {
255        self.visit_source(structure.loc)?;
256
257        Ok(())
258    }
259
260    fn visit_event(&mut self, event: &mut EventDefinition) -> Result<(), Self::Error> {
261        self.visit_source(event.loc)?;
262        self.visit_stray_semicolon()?;
263
264        Ok(())
265    }
266
267    fn visit_event_parameter(&mut self, param: &mut EventParameter) -> Result<(), Self::Error> {
268        self.visit_source(param.loc)
269    }
270
271    fn visit_error(&mut self, error: &mut ErrorDefinition) -> Result<(), Self::Error> {
272        self.visit_source(error.loc)?;
273        self.visit_stray_semicolon()?;
274
275        Ok(())
276    }
277
278    fn visit_error_parameter(&mut self, param: &mut ErrorParameter) -> Result<(), Self::Error> {
279        self.visit_source(param.loc)
280    }
281
282    fn visit_type_definition(&mut self, def: &mut TypeDefinition) -> Result<(), Self::Error> {
283        self.visit_source(def.loc)
284    }
285
286    fn visit_stray_semicolon(&mut self) -> Result<(), Self::Error> {
287        Ok(())
288    }
289
290    fn visit_opening_paren(&mut self) -> Result<(), Self::Error> {
291        Ok(())
292    }
293
294    fn visit_closing_paren(&mut self) -> Result<(), Self::Error> {
295        Ok(())
296    }
297
298    fn visit_newline(&mut self) -> Result<(), Self::Error> {
299        Ok(())
300    }
301
302    fn visit_using(&mut self, using: &mut Using) -> Result<(), Self::Error> {
303        self.visit_source(using.loc)?;
304        self.visit_stray_semicolon()?;
305
306        Ok(())
307    }
308
309    fn visit_yul_block(
310        &mut self,
311        loc: Loc,
312        _stmts: &mut Vec<YulStatement>,
313        _attempt_single_line: bool,
314    ) -> Result<(), Self::Error> {
315        self.visit_source(loc)
316    }
317
318    fn visit_yul_expr(&mut self, expr: &mut YulExpression) -> Result<(), Self::Error> {
319        self.visit_source(expr.loc())
320    }
321
322    fn visit_yul_assignment<T>(
323        &mut self,
324        loc: Loc,
325        _exprs: &mut Vec<T>,
326        _expr: &mut Option<&mut YulExpression>,
327    ) -> Result<(), Self::Error>
328    where
329        T: Visitable + CodeLocationExt,
330    {
331        self.visit_source(loc)
332    }
333
334    fn visit_yul_for(&mut self, stmt: &mut YulFor) -> Result<(), Self::Error> {
335        self.visit_source(stmt.loc)
336    }
337
338    fn visit_yul_function_call(&mut self, stmt: &mut YulFunctionCall) -> Result<(), Self::Error> {
339        self.visit_source(stmt.loc)
340    }
341
342    fn visit_yul_fun_def(&mut self, stmt: &mut YulFunctionDefinition) -> Result<(), Self::Error> {
343        self.visit_source(stmt.loc)
344    }
345
346    fn visit_yul_if(
347        &mut self,
348        loc: Loc,
349        _expr: &mut YulExpression,
350        _block: &mut YulBlock,
351    ) -> Result<(), Self::Error> {
352        self.visit_source(loc)
353    }
354
355    fn visit_yul_leave(&mut self, loc: Loc) -> Result<(), Self::Error> {
356        self.visit_source(loc)
357    }
358
359    fn visit_yul_switch(&mut self, stmt: &mut YulSwitch) -> Result<(), Self::Error> {
360        self.visit_source(stmt.loc)
361    }
362
363    fn visit_yul_var_declaration(
364        &mut self,
365        loc: Loc,
366        _idents: &mut Vec<YulTypedIdentifier>,
367        _expr: &mut Option<YulExpression>,
368    ) -> Result<(), Self::Error> {
369        self.visit_source(loc)
370    }
371
372    fn visit_yul_typed_ident(&mut self, ident: &mut YulTypedIdentifier) -> Result<(), Self::Error> {
373        self.visit_source(ident.loc)
374    }
375
376    fn visit_parser_error(&mut self, loc: Loc) -> Result<(), Self::Error> {
377        self.visit_source(loc)
378    }
379}
380
381/// Visitable trait for [`solang_parser::pt`] types.
382///
383/// All [`solang_parser::pt`] types, such as [Statement], should implement the [Visitable] trait
384/// that accepts a trait [Visitor] implementation, which has various callback handles for Solidity
385/// Parse Tree nodes.
386///
387/// We want to take a `&mut self` to be able to implement some advanced features in the future such
388/// as modifying the Parse Tree before formatting it.
389pub trait Visitable {
390    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
391    where
392        V: Visitor;
393}
394
395impl<T> Visitable for &mut T
396where
397    T: Visitable,
398{
399    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
400    where
401        V: Visitor,
402    {
403        T::visit(self, v)
404    }
405}
406
407impl<T> Visitable for Option<T>
408where
409    T: Visitable,
410{
411    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
412    where
413        V: Visitor,
414    {
415        if let Some(inner) = self.as_mut() {
416            inner.visit(v)
417        } else {
418            Ok(())
419        }
420    }
421}
422
423impl<T> Visitable for Box<T>
424where
425    T: Visitable,
426{
427    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
428    where
429        V: Visitor,
430    {
431        T::visit(self, v)
432    }
433}
434
435impl<T> Visitable for Vec<T>
436where
437    T: Visitable,
438{
439    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
440    where
441        V: Visitor,
442    {
443        for item in self.iter_mut() {
444            item.visit(v)?;
445        }
446        Ok(())
447    }
448}
449
450impl Visitable for SourceUnitPart {
451    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
452    where
453        V: Visitor,
454    {
455        match self {
456            Self::ContractDefinition(contract) => v.visit_contract(contract),
457            Self::PragmaDirective(pragma) => v.visit_pragma(pragma),
458            Self::ImportDirective(import) => import.visit(v),
459            Self::EnumDefinition(enumeration) => v.visit_enum(enumeration),
460            Self::StructDefinition(structure) => v.visit_struct(structure),
461            Self::EventDefinition(event) => v.visit_event(event),
462            Self::ErrorDefinition(error) => v.visit_error(error),
463            Self::FunctionDefinition(function) => v.visit_function(function),
464            Self::VariableDefinition(variable) => v.visit_var_definition(variable),
465            Self::TypeDefinition(def) => v.visit_type_definition(def),
466            Self::StraySemicolon(_) => v.visit_stray_semicolon(),
467            Self::Using(using) => v.visit_using(using),
468            Self::Annotation(annotation) => v.visit_annotation(annotation),
469        }
470    }
471}
472
473impl Visitable for Import {
474    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
475    where
476        V: Visitor,
477    {
478        match self {
479            Self::Plain(import, loc) => v.visit_import_plain(*loc, import),
480            Self::GlobalSymbol(global, import_as, loc) => {
481                v.visit_import_global(*loc, global, import_as)
482            }
483            Self::Rename(from, imports, loc) => v.visit_import_renames(*loc, imports, from),
484        }
485    }
486}
487
488impl Visitable for ContractPart {
489    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
490    where
491        V: Visitor,
492    {
493        match self {
494            Self::StructDefinition(structure) => v.visit_struct(structure),
495            Self::EventDefinition(event) => v.visit_event(event),
496            Self::ErrorDefinition(error) => v.visit_error(error),
497            Self::EnumDefinition(enumeration) => v.visit_enum(enumeration),
498            Self::VariableDefinition(variable) => v.visit_var_definition(variable),
499            Self::FunctionDefinition(function) => v.visit_function(function),
500            Self::TypeDefinition(def) => v.visit_type_definition(def),
501            Self::StraySemicolon(_) => v.visit_stray_semicolon(),
502            Self::Using(using) => v.visit_using(using),
503            Self::Annotation(annotation) => v.visit_annotation(annotation),
504        }
505    }
506}
507
508impl Visitable for Statement {
509    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
510    where
511        V: Visitor,
512    {
513        match self {
514            Self::Block { loc, unchecked, statements } => {
515                v.visit_block(*loc, *unchecked, statements)
516            }
517            Self::Assembly { loc, dialect, block, flags } => {
518                v.visit_assembly(*loc, dialect, block, flags)
519            }
520            Self::Args(loc, args) => v.visit_args(*loc, args),
521            Self::If(loc, cond, if_branch, else_branch) => {
522                v.visit_if(*loc, cond, if_branch, else_branch, true)
523            }
524            Self::While(loc, cond, body) => v.visit_while(*loc, cond, body),
525            Self::Expression(loc, expr) => {
526                v.visit_expr(*loc, expr)?;
527                v.visit_stray_semicolon()
528            }
529            Self::VariableDefinition(loc, declaration, expr) => {
530                v.visit_var_definition_stmt(*loc, declaration, expr)
531            }
532            Self::For(loc, init, cond, update, body) => v.visit_for(*loc, init, cond, update, body),
533            Self::DoWhile(loc, body, cond) => v.visit_do_while(*loc, body, cond),
534            Self::Continue(loc) => v.visit_continue(*loc, true),
535            Self::Break(loc) => v.visit_break(*loc, true),
536            Self::Return(loc, expr) => v.visit_return(*loc, expr),
537            Self::Revert(loc, error, args) => v.visit_revert(*loc, error, args),
538            Self::RevertNamedArgs(loc, error, args) => v.visit_revert_named_args(*loc, error, args),
539            Self::Emit(loc, event) => v.visit_emit(*loc, event),
540            Self::Try(loc, expr, returns, clauses) => v.visit_try(*loc, expr, returns, clauses),
541            Self::Error(loc) => v.visit_parser_error(*loc),
542        }
543    }
544}
545
546impl Visitable for Loc {
547    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
548    where
549        V: Visitor,
550    {
551        v.visit_source(*self)
552    }
553}
554
555impl Visitable for Expression {
556    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
557    where
558        V: Visitor,
559    {
560        v.visit_expr(self.loc(), self)
561    }
562}
563
564impl Visitable for Identifier {
565    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
566    where
567        V: Visitor,
568    {
569        v.visit_ident(self.loc, self)
570    }
571}
572
573impl Visitable for VariableDeclaration {
574    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
575    where
576        V: Visitor,
577    {
578        v.visit_var_declaration(self)
579    }
580}
581
582impl Visitable for YulBlock {
583    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
584    where
585        V: Visitor,
586    {
587        v.visit_yul_block(self.loc, self.statements.as_mut(), false)
588    }
589}
590
591impl Visitable for YulStatement {
592    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
593    where
594        V: Visitor,
595    {
596        match self {
597            Self::Assign(loc, exprs, expr) => v.visit_yul_assignment(*loc, exprs, &mut Some(expr)),
598            Self::Block(block) => v.visit_yul_block(block.loc, block.statements.as_mut(), false),
599            Self::Break(loc) => v.visit_break(*loc, false),
600            Self::Continue(loc) => v.visit_continue(*loc, false),
601            Self::For(stmt) => v.visit_yul_for(stmt),
602            Self::FunctionCall(stmt) => v.visit_yul_function_call(stmt),
603            Self::FunctionDefinition(stmt) => v.visit_yul_fun_def(stmt),
604            Self::If(loc, expr, block) => v.visit_yul_if(*loc, expr, block),
605            Self::Leave(loc) => v.visit_yul_leave(*loc),
606            Self::Switch(stmt) => v.visit_yul_switch(stmt),
607            Self::VariableDeclaration(loc, idents, expr) => {
608                v.visit_yul_var_declaration(*loc, idents, expr)
609            }
610            Self::Error(loc) => v.visit_parser_error(*loc),
611        }
612    }
613}
614
615macro_rules! impl_visitable {
616    ($type:ty, $func:ident) => {
617        impl Visitable for $type {
618            fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
619            where
620                V: Visitor,
621            {
622                v.$func(self)
623            }
624        }
625    };
626}
627
628impl_visitable!(SourceUnit, visit_source_unit);
629impl_visitable!(FunctionAttribute, visit_function_attribute);
630impl_visitable!(VariableAttribute, visit_var_attribute);
631impl_visitable!(Parameter, visit_parameter);
632impl_visitable!(Base, visit_base);
633impl_visitable!(EventParameter, visit_event_parameter);
634impl_visitable!(ErrorParameter, visit_error_parameter);
635impl_visitable!(IdentifierPath, visit_ident_path);
636impl_visitable!(YulExpression, visit_yul_expr);
637impl_visitable!(YulTypedIdentifier, visit_yul_typed_ident);