forge_fmt/
visit.rs

1//! Visitor helpers to traverse the [solang Solidity Parse Tree](solang_parser::pt).
2
3use crate::solang_ext::{CodeLocationExt, pt::*};
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() { inner.visit(v) } else { Ok(()) }
416    }
417}
418
419impl<T> Visitable for Box<T>
420where
421    T: Visitable,
422{
423    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
424    where
425        V: Visitor,
426    {
427        T::visit(self, v)
428    }
429}
430
431impl<T> Visitable for Vec<T>
432where
433    T: Visitable,
434{
435    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
436    where
437        V: Visitor,
438    {
439        for item in self.iter_mut() {
440            item.visit(v)?;
441        }
442        Ok(())
443    }
444}
445
446impl Visitable for SourceUnitPart {
447    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
448    where
449        V: Visitor,
450    {
451        match self {
452            Self::ContractDefinition(contract) => v.visit_contract(contract),
453            Self::PragmaDirective(pragma) => v.visit_pragma(pragma),
454            Self::ImportDirective(import) => import.visit(v),
455            Self::EnumDefinition(enumeration) => v.visit_enum(enumeration),
456            Self::StructDefinition(structure) => v.visit_struct(structure),
457            Self::EventDefinition(event) => v.visit_event(event),
458            Self::ErrorDefinition(error) => v.visit_error(error),
459            Self::FunctionDefinition(function) => v.visit_function(function),
460            Self::VariableDefinition(variable) => v.visit_var_definition(variable),
461            Self::TypeDefinition(def) => v.visit_type_definition(def),
462            Self::StraySemicolon(_) => v.visit_stray_semicolon(),
463            Self::Using(using) => v.visit_using(using),
464            Self::Annotation(annotation) => v.visit_annotation(annotation),
465        }
466    }
467}
468
469impl Visitable for Import {
470    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
471    where
472        V: Visitor,
473    {
474        match self {
475            Self::Plain(import, loc) => v.visit_import_plain(*loc, import),
476            Self::GlobalSymbol(global, import_as, loc) => {
477                v.visit_import_global(*loc, global, import_as)
478            }
479            Self::Rename(from, imports, loc) => v.visit_import_renames(*loc, imports, from),
480        }
481    }
482}
483
484impl Visitable for ContractPart {
485    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
486    where
487        V: Visitor,
488    {
489        match self {
490            Self::StructDefinition(structure) => v.visit_struct(structure),
491            Self::EventDefinition(event) => v.visit_event(event),
492            Self::ErrorDefinition(error) => v.visit_error(error),
493            Self::EnumDefinition(enumeration) => v.visit_enum(enumeration),
494            Self::VariableDefinition(variable) => v.visit_var_definition(variable),
495            Self::FunctionDefinition(function) => v.visit_function(function),
496            Self::TypeDefinition(def) => v.visit_type_definition(def),
497            Self::StraySemicolon(_) => v.visit_stray_semicolon(),
498            Self::Using(using) => v.visit_using(using),
499            Self::Annotation(annotation) => v.visit_annotation(annotation),
500        }
501    }
502}
503
504impl Visitable for Statement {
505    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
506    where
507        V: Visitor,
508    {
509        match self {
510            Self::Block { loc, unchecked, statements } => {
511                v.visit_block(*loc, *unchecked, statements)
512            }
513            Self::Assembly { loc, dialect, block, flags } => {
514                v.visit_assembly(*loc, dialect, block, flags)
515            }
516            Self::Args(loc, args) => v.visit_args(*loc, args),
517            Self::If(loc, cond, if_branch, else_branch) => {
518                v.visit_if(*loc, cond, if_branch, else_branch, true)
519            }
520            Self::While(loc, cond, body) => v.visit_while(*loc, cond, body),
521            Self::Expression(loc, expr) => {
522                v.visit_expr(*loc, expr)?;
523                v.visit_stray_semicolon()
524            }
525            Self::VariableDefinition(loc, declaration, expr) => {
526                v.visit_var_definition_stmt(*loc, declaration, expr)
527            }
528            Self::For(loc, init, cond, update, body) => v.visit_for(*loc, init, cond, update, body),
529            Self::DoWhile(loc, body, cond) => v.visit_do_while(*loc, body, cond),
530            Self::Continue(loc) => v.visit_continue(*loc, true),
531            Self::Break(loc) => v.visit_break(*loc, true),
532            Self::Return(loc, expr) => v.visit_return(*loc, expr),
533            Self::Revert(loc, error, args) => v.visit_revert(*loc, error, args),
534            Self::RevertNamedArgs(loc, error, args) => v.visit_revert_named_args(*loc, error, args),
535            Self::Emit(loc, event) => v.visit_emit(*loc, event),
536            Self::Try(loc, expr, returns, clauses) => v.visit_try(*loc, expr, returns, clauses),
537            Self::Error(loc) => v.visit_parser_error(*loc),
538        }
539    }
540}
541
542impl Visitable for Loc {
543    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
544    where
545        V: Visitor,
546    {
547        v.visit_source(*self)
548    }
549}
550
551impl Visitable for Expression {
552    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
553    where
554        V: Visitor,
555    {
556        v.visit_expr(self.loc(), self)
557    }
558}
559
560impl Visitable for Identifier {
561    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
562    where
563        V: Visitor,
564    {
565        v.visit_ident(self.loc, self)
566    }
567}
568
569impl Visitable for VariableDeclaration {
570    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
571    where
572        V: Visitor,
573    {
574        v.visit_var_declaration(self)
575    }
576}
577
578impl Visitable for YulBlock {
579    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
580    where
581        V: Visitor,
582    {
583        v.visit_yul_block(self.loc, self.statements.as_mut(), false)
584    }
585}
586
587impl Visitable for YulStatement {
588    fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
589    where
590        V: Visitor,
591    {
592        match self {
593            Self::Assign(loc, exprs, expr) => v.visit_yul_assignment(*loc, exprs, &mut Some(expr)),
594            Self::Block(block) => v.visit_yul_block(block.loc, block.statements.as_mut(), false),
595            Self::Break(loc) => v.visit_break(*loc, false),
596            Self::Continue(loc) => v.visit_continue(*loc, false),
597            Self::For(stmt) => v.visit_yul_for(stmt),
598            Self::FunctionCall(stmt) => v.visit_yul_function_call(stmt),
599            Self::FunctionDefinition(stmt) => v.visit_yul_fun_def(stmt),
600            Self::If(loc, expr, block) => v.visit_yul_if(*loc, expr, block),
601            Self::Leave(loc) => v.visit_yul_leave(*loc),
602            Self::Switch(stmt) => v.visit_yul_switch(stmt),
603            Self::VariableDeclaration(loc, idents, expr) => {
604                v.visit_yul_var_declaration(*loc, idents, expr)
605            }
606            Self::Error(loc) => v.visit_parser_error(*loc),
607        }
608    }
609}
610
611macro_rules! impl_visitable {
612    ($type:ty, $func:ident) => {
613        impl Visitable for $type {
614            fn visit<V>(&mut self, v: &mut V) -> Result<(), V::Error>
615            where
616                V: Visitor,
617            {
618                v.$func(self)
619            }
620        }
621    };
622}
623
624impl_visitable!(SourceUnit, visit_source_unit);
625impl_visitable!(FunctionAttribute, visit_function_attribute);
626impl_visitable!(VariableAttribute, visit_var_attribute);
627impl_visitable!(Parameter, visit_parameter);
628impl_visitable!(Base, visit_base);
629impl_visitable!(EventParameter, visit_event_parameter);
630impl_visitable!(ErrorParameter, visit_error_parameter);
631impl_visitable!(IdentifierPath, visit_ident_path);
632impl_visitable!(YulExpression, visit_yul_expr);
633impl_visitable!(YulTypedIdentifier, visit_yul_typed_ident);