Skip to main content

plc_ast/
mut_visitor.rs

1//! This module defines the `AstVisitorMut` trait and its associated macros.
2//! The `AstVisitorMut` trait provides a set of methods for mutably traversing and visiting ASTs
3
4use std::borrow::BorrowMut;
5
6use crate::ast::{
7    flatten_expression_list, Assignment, AstNode, AstStatement, BinaryExpression, CallStatement,
8    CompilationUnit, DataType, DataTypeDeclaration, DirectAccess, HardwareAccess, Implementation, Interface,
9    JumpStatement, MultipliedStatement, Pou, PropertyBlock, RangeStatement, ReferenceAccess, ReferenceExpr,
10    UnaryExpression, UserTypeDeclaration, Variable, VariableBlock,
11};
12use crate::control_statements::{
13    AstControlStatement, ConditionalBlock, ForLoopStatement, LoopStatement, ReturnStatement,
14};
15use crate::literals::AstLiteral;
16use crate::try_from_mut;
17
18#[macro_export]
19macro_rules! visit_all_nodes_mut {
20    ($visitor:expr, $iter:expr) => {
21        // Note: The `allow` is needed to suppress warnings about `while let Some(...)` warnings
22        // because `visit_all_nodes_mut!` is used for both Option and Non-Option types
23        #[allow(warnings)]
24        {
25            for node in $iter {
26                $visitor.visit(node);
27            }
28        }
29    };
30}
31
32/// Macro that calls the visitor's `visit` method for every AstNode in the passed sequence of nodes.
33macro_rules! visit_nodes_mut {
34    ($visitor:expr, $($node:expr),*) => {
35        $(
36            $visitor.visit($node);
37        )*
38    };
39}
40
41pub trait WalkerMut {
42    fn walk<V>(&mut self, visitor: &mut V)
43    where
44        V: AstVisitorMut;
45}
46
47pub trait AstVisitorMut: Sized {
48    fn visit(&mut self, node: &mut AstNode) {
49        node.walk(self)
50    }
51
52    /// Called when visiting a list of statements (e.g. implementation bodies, control flow branches).
53    /// Override this to intercept statement-list processing — for example, to drain-and-rebuild the
54    /// list for 1→N statement expansion.
55    fn visit_statement_list(&mut self, stmts: &mut Vec<AstNode>) {
56        for node in stmts.iter_mut() {
57            self.visit(node);
58        }
59    }
60
61    //Takes ownership of the node, manipulates it and returns a new node
62    fn map(&mut self, mut node: AstNode) -> AstNode {
63        node.borrow_mut().walk(self);
64        node
65    }
66
67    fn visit_compilation_unit(&mut self, unit: &mut CompilationUnit) {
68        unit.walk(self)
69    }
70
71    fn visit_implementation(&mut self, implementation: &mut Implementation) {
72        implementation.walk(self);
73    }
74
75    fn visit_variable_block(&mut self, block: &mut VariableBlock) {
76        block.walk(self)
77    }
78
79    fn visit_variable(&mut self, variable: &mut Variable) {
80        variable.walk(self);
81    }
82
83    fn visit_enum_element(&mut self, element: &mut AstNode) {
84        element.walk(self);
85    }
86
87    fn visit_data_type_declaration(&mut self, data_type_declaration: &mut DataTypeDeclaration) {
88        data_type_declaration.walk(self);
89    }
90
91    fn visit_user_type_declaration(&mut self, user_type: &mut UserTypeDeclaration) {
92        user_type.walk(self);
93    }
94
95    fn visit_data_type(&mut self, data_type: &mut DataType) {
96        data_type.walk(self);
97    }
98
99    fn visit_pou(&mut self, pou: &mut Pou) {
100        pou.walk(self);
101    }
102
103    fn visit_empty_statement(&mut self, _node: &mut AstNode) {}
104
105    fn visit_default_value(&mut self, _node: &mut AstNode) {}
106
107    fn visit_literal(&mut self, node: &mut AstNode) {
108        let stmt = try_from_mut!(node, AstLiteral).expect("Is a literal");
109        stmt.walk(self)
110    }
111
112    fn visit_multiplied_statement(&mut self, node: &mut AstNode) {
113        let stmt = try_from_mut!(node, MultipliedStatement).expect("MultipliedStatement");
114        stmt.walk(self)
115    }
116
117    fn visit_reference_expr(&mut self, node: &mut AstNode) {
118        let stmt = try_from_mut!(node, ReferenceExpr).expect("ReferenceExpr");
119        stmt.walk(self)
120    }
121
122    fn visit_identifier(&mut self, _node: &mut AstNode) {}
123
124    fn visit_direct_access(&mut self, node: &mut AstNode) {
125        let stmt = try_from_mut!(node, DirectAccess).expect("DirectAccess");
126        stmt.walk(self)
127    }
128
129    fn visit_hardware_access(&mut self, node: &mut AstNode) {
130        let stmt = try_from_mut!(node, HardwareAccess).expect("HardwareAccess");
131        stmt.walk(self)
132    }
133
134    fn visit_binary_expression(&mut self, node: &mut AstNode) {
135        let stmt = try_from_mut!(node, BinaryExpression).expect("BinaryExpression");
136        stmt.walk(self)
137    }
138
139    fn visit_unary_expression(&mut self, node: &mut AstNode) {
140        let stmt = try_from_mut!(node, UnaryExpression).expect("UnaryExpression");
141        stmt.walk(self)
142    }
143
144    fn visit_expression_list(&mut self, node: &mut AstNode) {
145        let stmt = try_from_mut!(node, Vec<AstNode>).expect("Vec<AstNode>");
146        visit_all_nodes_mut!(self, stmt);
147    }
148
149    fn visit_paren_expression(&mut self, node: &mut AstNode) {
150        let AstStatement::ParenExpression(inner) = node.get_stmt_mut() else {
151            unreachable!("Must be ParenExpression");
152        };
153        inner.walk(self)
154    }
155
156    fn visit_range_statement(&mut self, node: &mut AstNode) {
157        let stmt = try_from_mut!(node, RangeStatement).expect("RangeStatement");
158        stmt.walk(self)
159    }
160
161    fn visit_vla_range_statement(&mut self, _node: &mut AstNode) {}
162
163    fn visit_assignment(&mut self, node: &mut AstNode) {
164        let stmt = try_from_mut!(node, Assignment).expect("Assignment");
165        stmt.walk(self)
166    }
167
168    fn visit_output_assignment(&mut self, node: &mut AstNode) {
169        let stmt = try_from_mut!(node, Assignment).expect("Assignment");
170        stmt.walk(self)
171    }
172
173    fn visit_ref_assignment(&mut self, node: &mut AstNode) {
174        let stmt = try_from_mut!(node, Assignment).expect("Assignment");
175        stmt.walk(self)
176    }
177
178    fn visit_call_statement(&mut self, node: &mut AstNode) {
179        let stmt = try_from_mut!(node, CallStatement).expect("CallStatement");
180        stmt.walk(self)
181    }
182
183    fn visit_control_statement(&mut self, node: &mut AstNode) {
184        let stmt = try_from_mut!(node, AstControlStatement).expect("AstControlStatement");
185        match stmt {
186            AstControlStatement::ForLoop(for_stmt) => self.visit_for_loop_statement(for_stmt),
187            AstControlStatement::WhileLoop(loop_stmt) => self.visit_while_loop_statement(loop_stmt),
188            AstControlStatement::RepeatLoop(loop_stmt) => self.visit_repeat_loop_statement(loop_stmt),
189            _ => stmt.walk(self),
190        }
191    }
192
193    /// Visits a `ForLoop` control statement.
194    /// Make sure to visit the counter, bounds, optional step and body to continue the traversal.
195    /// # Arguments
196    /// * `stmt` - The unwraped, typed `ForLoopStatement` node to visit.
197    fn visit_for_loop_statement(&mut self, stmt: &mut ForLoopStatement) {
198        visit_nodes_mut!(self, &mut stmt.counter, &mut stmt.start, &mut stmt.end);
199        visit_all_nodes_mut!(self, &mut stmt.by_step);
200        self.visit_statement_list(&mut stmt.body);
201    }
202
203    /// Visits a `WhileLoop` control statement.
204    /// Make sure to visit the condition and body to continue the traversal.
205    /// # Arguments
206    /// * `stmt` - The unwraped, typed `LoopStatement` node to visit.
207    fn visit_while_loop_statement(&mut self, stmt: &mut LoopStatement) {
208        visit_nodes_mut!(self, &mut stmt.condition);
209        self.visit_statement_list(&mut stmt.body);
210    }
211
212    /// Visits a `RepeatLoop` control statement.
213    /// Make sure to visit the condition and body to continue the traversal.
214    /// # Arguments
215    /// * `stmt` - The unwraped, typed `LoopStatement` node to visit.
216    fn visit_repeat_loop_statement(&mut self, stmt: &mut LoopStatement) {
217        visit_nodes_mut!(self, &mut stmt.condition);
218        self.visit_statement_list(&mut stmt.body);
219    }
220
221    fn visit_case_condition(&mut self, node: &mut AstNode) {
222        let AstStatement::CaseCondition(child) = node.get_stmt_mut() else {
223            unreachable!("CaseCondition");
224        };
225        child.walk(self)
226    }
227
228    fn visit_exit_statement(&mut self, _node: &mut AstNode) {}
229
230    fn visit_continue_statement(&mut self, _node: &mut AstNode) {}
231
232    fn visit_return_statement(&mut self, node: &mut AstNode) {
233        let stmt = try_from_mut!(node, ReturnStatement).expect("ReturnStatement");
234        stmt.walk(self)
235    }
236
237    fn visit_jump_statement(&mut self, node: &mut AstNode) {
238        let stmt = try_from_mut!(node, JumpStatement).expect("CallStatement");
239        stmt.walk(self)
240    }
241
242    /// Visits a `LabelStatement` node.
243    /// # Arguments
244    /// * `stmt` - The unwrapedyped `LabelStatement` node to visit.
245    /// * `node` - The wrapped `AstNode` node to visit. Offers access to location information and AstId
246    fn visit_label_statement(&mut self, _node: &mut AstNode) {}
247
248    /// Visits a `Allocation` node.
249    /// # Arguments
250    /// * `stmt` - The unwrapedyped `Allocation` node to visit.
251    /// * `node` - The wrapped `AstNode` node to visit. Offers access to location information and AstId
252    fn visit_allocation(&mut self, _node: &mut AstNode) {}
253
254    fn visit_interface(&mut self, interface: &mut Interface) {
255        interface.walk(self);
256    }
257
258    fn visit_property(&mut self, property: &mut PropertyBlock) {
259        property.walk(self);
260    }
261
262    fn visit_super(&mut self, _node: &mut AstNode) {}
263
264    fn visit_this(&mut self, _node: &mut AstNode) {}
265}
266
267impl WalkerMut for AstLiteral {
268    fn walk<V>(&mut self, _visitor: &mut V)
269    where
270        V: AstVisitorMut,
271    {
272        // do nothing
273    }
274}
275
276impl WalkerMut for MultipliedStatement {
277    fn walk<V>(&mut self, visitor: &mut V)
278    where
279        V: AstVisitorMut,
280    {
281        visitor.visit(&mut self.element)
282    }
283}
284
285impl WalkerMut for ReferenceExpr {
286    fn walk<V>(&mut self, visitor: &mut V)
287    where
288        V: AstVisitorMut,
289    {
290        if let Some(base) = &mut self.base {
291            visitor.visit(base);
292        }
293
294        match &mut self.access {
295            ReferenceAccess::Member(t) | ReferenceAccess::Index(t) | ReferenceAccess::Cast(t) => {
296                visitor.visit(t)
297            }
298            _ => {}
299        }
300    }
301}
302
303impl WalkerMut for DirectAccess {
304    fn walk<V>(&mut self, visitor: &mut V)
305    where
306        V: AstVisitorMut,
307    {
308        visit_nodes_mut!(visitor, &mut self.index);
309    }
310}
311
312impl WalkerMut for HardwareAccess {
313    fn walk<V>(&mut self, visitor: &mut V)
314    where
315        V: AstVisitorMut,
316    {
317        visit_all_nodes_mut!(visitor, &mut self.address);
318    }
319}
320
321impl WalkerMut for BinaryExpression {
322    fn walk<V>(&mut self, visitor: &mut V)
323    where
324        V: AstVisitorMut,
325    {
326        visit_nodes_mut!(visitor, &mut self.left, &mut self.right);
327    }
328}
329
330impl WalkerMut for UnaryExpression {
331    fn walk<V>(&mut self, visitor: &mut V)
332    where
333        V: AstVisitorMut,
334    {
335        visit_nodes_mut!(visitor, &mut self.value);
336    }
337}
338
339impl WalkerMut for Assignment {
340    fn walk<V>(&mut self, visitor: &mut V)
341    where
342        V: AstVisitorMut,
343    {
344        visit_nodes_mut!(visitor, &mut self.left, &mut self.right);
345    }
346}
347
348impl WalkerMut for RangeStatement {
349    fn walk<V>(&mut self, visitor: &mut V)
350    where
351        V: AstVisitorMut,
352    {
353        visit_nodes_mut!(visitor, &mut self.start, &mut self.end);
354    }
355}
356
357impl WalkerMut for CallStatement {
358    fn walk<V>(&mut self, visitor: &mut V)
359    where
360        V: AstVisitorMut,
361    {
362        visit_nodes_mut!(visitor, &mut self.operator);
363        if let Some(params) = &mut self.parameters {
364            visit_nodes_mut!(visitor, params);
365        }
366    }
367}
368
369impl WalkerMut for Vec<ConditionalBlock> {
370    fn walk<V>(&mut self, visitor: &mut V)
371    where
372        V: AstVisitorMut,
373    {
374        for b in self {
375            visit_nodes_mut!(visitor, &mut b.condition);
376            visitor.visit_statement_list(&mut b.body);
377        }
378    }
379}
380
381impl WalkerMut for AstControlStatement {
382    fn walk<V>(&mut self, visitor: &mut V)
383    where
384        V: AstVisitorMut,
385    {
386        match self {
387            AstControlStatement::If(stmt) => {
388                stmt.blocks.walk(visitor);
389                visitor.visit_statement_list(&mut stmt.else_block);
390            }
391            AstControlStatement::WhileLoop(stmt) | AstControlStatement::RepeatLoop(stmt) => {
392                visit_nodes_mut!(visitor, &mut stmt.condition);
393                visitor.visit_statement_list(&mut stmt.body);
394            }
395            AstControlStatement::ForLoop(stmt) => {
396                visit_nodes_mut!(visitor, &mut stmt.counter, &mut stmt.start, &mut stmt.end);
397                visit_all_nodes_mut!(visitor, &mut stmt.by_step);
398                visitor.visit_statement_list(&mut stmt.body);
399            }
400            AstControlStatement::Case(stmt) => {
401                visit_nodes_mut!(visitor, &mut stmt.selector);
402                stmt.case_blocks.walk(visitor);
403                visitor.visit_statement_list(&mut stmt.else_block);
404            }
405        }
406    }
407}
408
409impl WalkerMut for ReturnStatement {
410    fn walk<V>(&mut self, visitor: &mut V)
411    where
412        V: AstVisitorMut,
413    {
414        visit_all_nodes_mut!(visitor, &mut self.condition);
415    }
416}
417
418impl WalkerMut for JumpStatement {
419    fn walk<V>(&mut self, visitor: &mut V)
420    where
421        V: AstVisitorMut,
422    {
423        visit_nodes_mut!(visitor, &mut self.condition, &mut self.target);
424    }
425}
426
427impl WalkerMut for Interface {
428    fn walk<V>(&mut self, visitor: &mut V)
429    where
430        V: AstVisitorMut,
431    {
432        for method in &mut self.methods {
433            visitor.visit_pou(method);
434        }
435
436        for property in &mut self.properties {
437            visitor.visit_property(property);
438        }
439    }
440}
441
442impl WalkerMut for PropertyBlock {
443    fn walk<V>(&mut self, visitor: &mut V)
444    where
445        V: AstVisitorMut,
446    {
447        for implementation in &mut self.implementations {
448            visitor.visit_data_type_declaration(&mut implementation.datatype);
449            for block in &mut implementation.variable_blocks {
450                visitor.visit_variable_block(block);
451            }
452            visitor.visit_statement_list(&mut implementation.body);
453        }
454    }
455}
456
457impl WalkerMut for AstNode {
458    fn walk<V>(&mut self, visitor: &mut V)
459    where
460        V: AstVisitorMut,
461    {
462        match self.stmt {
463            AstStatement::EmptyStatement(_) => visitor.visit_empty_statement(self),
464            AstStatement::DefaultValue(_) => visitor.visit_default_value(self),
465            AstStatement::Literal(_) => visitor.visit_literal(self),
466            AstStatement::MultipliedStatement(_) => visitor.visit_multiplied_statement(self),
467            AstStatement::ReferenceExpr(_) => visitor.visit_reference_expr(self),
468            AstStatement::Identifier(_) => visitor.visit_identifier(self),
469            AstStatement::DirectAccess(_) => visitor.visit_direct_access(self),
470            AstStatement::HardwareAccess(_) => visitor.visit_hardware_access(self),
471            AstStatement::BinaryExpression(_) => visitor.visit_binary_expression(self),
472            AstStatement::UnaryExpression(_) => visitor.visit_unary_expression(self),
473            AstStatement::ExpressionList(_) => visitor.visit_expression_list(self),
474            AstStatement::ParenExpression(_) => visitor.visit_paren_expression(self),
475            AstStatement::RangeStatement(_) => visitor.visit_range_statement(self),
476            AstStatement::VlaRangeStatement => visitor.visit_vla_range_statement(self),
477            AstStatement::Assignment(_) => visitor.visit_assignment(self),
478            AstStatement::OutputAssignment(_) => visitor.visit_output_assignment(self),
479            AstStatement::RefAssignment(_) => visitor.visit_ref_assignment(self),
480            AstStatement::CallStatement(_) => visitor.visit_call_statement(self),
481            AstStatement::ControlStatement(_) => visitor.visit_control_statement(self),
482            AstStatement::CaseCondition(_) => visitor.visit_case_condition(self),
483            AstStatement::ExitStatement(_) => visitor.visit_exit_statement(self),
484            AstStatement::ContinueStatement(_) => visitor.visit_continue_statement(self),
485            AstStatement::ReturnStatement(_) => visitor.visit_return_statement(self),
486            AstStatement::JumpStatement(_) => visitor.visit_jump_statement(self),
487            AstStatement::LabelStatement(_) => visitor.visit_label_statement(self),
488            AstStatement::AllocationStatement(_) => visitor.visit_allocation(self),
489            AstStatement::Super(_) => visitor.visit_super(self),
490            AstStatement::This => visitor.visit_this(self),
491        }
492    }
493}
494
495impl WalkerMut for CompilationUnit {
496    fn walk<V>(&mut self, visitor: &mut V)
497    where
498        V: AstVisitorMut,
499    {
500        for block in &mut self.global_vars {
501            visitor.visit_variable_block(block);
502        }
503
504        for user_type in &mut self.user_types {
505            visitor.visit_user_type_declaration(user_type);
506        }
507
508        for pou in &mut self.pous {
509            visitor.visit_pou(pou);
510        }
511
512        for i in &mut self.implementations {
513            visitor.visit_implementation(i);
514        }
515
516        for i in &mut self.interfaces {
517            visitor.visit_interface(i);
518        }
519    }
520}
521
522impl WalkerMut for UserTypeDeclaration {
523    fn walk<V>(&mut self, visitor: &mut V)
524    where
525        V: AstVisitorMut,
526    {
527        visitor.visit_data_type(&mut self.data_type);
528        visit_all_nodes_mut!(visitor, &mut self.initializer);
529    }
530}
531
532impl WalkerMut for VariableBlock {
533    fn walk<V>(&mut self, visitor: &mut V)
534    where
535        V: AstVisitorMut,
536    {
537        for v in self.variables.iter_mut() {
538            visitor.visit_variable(v);
539        }
540    }
541}
542
543impl WalkerMut for Variable {
544    fn walk<V>(&mut self, visitor: &mut V)
545    where
546        V: AstVisitorMut,
547    {
548        visit_all_nodes_mut!(visitor, &mut self.address);
549        visitor.visit_data_type_declaration(&mut self.data_type_declaration);
550        visit_all_nodes_mut!(visitor, &mut self.initializer);
551    }
552}
553
554impl WalkerMut for DataType {
555    fn walk<V>(&mut self, visitor: &mut V)
556    where
557        V: AstVisitorMut,
558    {
559        match self {
560            DataType::StructType { variables, .. } => {
561                for v in variables.iter_mut() {
562                    visitor.visit_variable(v);
563                }
564            }
565            DataType::EnumType { elements, .. } => {
566                flatten_expression_list(elements).iter_mut().map(|it| it.clone()).for_each(|mut ele| {
567                    visitor.visit_enum_element(&mut ele);
568                });
569            }
570            DataType::SubRangeType { bounds, .. } => {
571                visit_all_nodes_mut!(visitor, bounds);
572            }
573            DataType::ArrayType { bounds, referenced_type, .. } => {
574                visitor.visit(bounds);
575                visitor.visit_data_type_declaration(referenced_type);
576            }
577            DataType::PointerType { referenced_type, .. } => {
578                visitor.visit_data_type_declaration(referenced_type);
579            }
580            DataType::StringType { size, .. } => {
581                visit_all_nodes_mut!(visitor, size);
582            }
583            DataType::VarArgs { referenced_type, .. } => {
584                if let Some(data_type_declaration) = referenced_type {
585                    visitor.visit_data_type_declaration(data_type_declaration);
586                }
587            }
588            DataType::GenericType { .. } => {
589                //no further visits
590            }
591        }
592    }
593}
594
595impl WalkerMut for DataTypeDeclaration {
596    fn walk<V>(&mut self, visitor: &mut V)
597    where
598        V: AstVisitorMut,
599    {
600        if let DataTypeDeclaration::Definition { data_type, .. } = self {
601            visitor.visit_data_type(data_type);
602        }
603    }
604}
605
606impl<T: WalkerMut> WalkerMut for Option<T> {
607    fn walk<V>(&mut self, visitor: &mut V)
608    where
609        V: AstVisitorMut,
610    {
611        if let Some(node) = self {
612            node.walk(visitor);
613        }
614    }
615}
616
617impl WalkerMut for Pou {
618    fn walk<V>(&mut self, visitor: &mut V)
619    where
620        V: AstVisitorMut,
621    {
622        for block in &mut self.variable_blocks {
623            visitor.visit_variable_block(block);
624        }
625
626        for property in &mut self.properties {
627            visitor.visit_property(property);
628        }
629
630        if let Some(rt) = self.return_type.as_mut() {
631            visitor.visit_data_type_declaration(rt)
632        }
633    }
634}
635
636impl WalkerMut for Implementation {
637    fn walk<V>(&mut self, visitor: &mut V)
638    where
639        V: AstVisitorMut,
640    {
641        visitor.visit_statement_list(&mut self.statements);
642    }
643}