1use 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 #[allow(warnings)]
24 {
25 for node in $iter {
26 $visitor.visit(node);
27 }
28 }
29 };
30}
31
32macro_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 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 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 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 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 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 fn visit_label_statement(&mut self, _node: &mut AstNode) {}
247
248 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 }
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 }
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}