1use crate::ast::{Expression, FormatStringPart, MatchArm, SelectArm, SelectArmPattern};
2
3pub trait AstFolder {
4 type Error;
5
6 fn fold_module(
7 &mut self,
8 expressions: Vec<Expression>,
9 ) -> Result<Vec<Expression>, Self::Error> {
10 expressions
11 .into_iter()
12 .map(|e| self.fold_expression(e))
13 .collect()
14 }
15
16 fn fold_expression(&mut self, expression: Expression) -> Result<Expression, Self::Error> {
17 self.fold_expression_default(expression)
18 }
19
20 fn fold_expression_default(
21 &mut self,
22 expression: Expression,
23 ) -> Result<Expression, Self::Error> {
24 use Expression::*;
25
26 Ok(match expression {
27 Binary {
28 operator,
29 left,
30 right,
31 ty,
32 span,
33 } => Binary {
34 operator,
35 left: Box::new(self.fold_expression(*left)?),
36 right: Box::new(self.fold_expression(*right)?),
37 ty,
38 span,
39 },
40
41 Call {
42 expression,
43 args,
44 spread,
45 type_args,
46 ty,
47 span,
48 call_kind,
49 } => Call {
50 expression: Box::new(self.fold_expression(*expression)?),
51 args: self.fold_vec(args)?,
52 spread: Box::new((*spread).map(|e| self.fold_expression(e)).transpose()?),
53 type_args,
54 ty,
55 span,
56 call_kind,
57 },
58
59 Block { items, ty, span } => Block {
60 items: self.fold_vec(items)?,
61 ty,
62 span,
63 },
64
65 TryBlock {
66 items,
67 ty,
68 try_keyword_span,
69 span,
70 } => TryBlock {
71 items: self.fold_vec(items)?,
72 ty,
73 try_keyword_span,
74 span,
75 },
76
77 RecoverBlock {
78 items,
79 ty,
80 recover_keyword_span,
81 span,
82 } => RecoverBlock {
83 items: self.fold_vec(items)?,
84 ty,
85 recover_keyword_span,
86 span,
87 },
88
89 If {
90 condition,
91 consequence,
92 alternative,
93 ty,
94 span,
95 } => If {
96 condition: Box::new(self.fold_expression(*condition)?),
97 consequence: Box::new(self.fold_expression(*consequence)?),
98 alternative: Box::new(self.fold_expression(*alternative)?),
99 ty,
100 span,
101 },
102
103 IfLet {
104 pattern,
105 scrutinee,
106 consequence,
107 alternative,
108 typed_pattern,
109 else_span,
110 ty,
111 span,
112 } => IfLet {
113 pattern,
114 scrutinee: Box::new(self.fold_expression(*scrutinee)?),
115 consequence: Box::new(self.fold_expression(*consequence)?),
116 alternative: Box::new(self.fold_expression(*alternative)?),
117 typed_pattern,
118 else_span,
119 ty,
120 span,
121 },
122
123 Match {
124 subject,
125 arms,
126 origin,
127 ty,
128 span,
129 } => Match {
130 subject: Box::new(self.fold_expression(*subject)?),
131 arms: arms
132 .into_iter()
133 .map(|arm| self.fold_match_arm(arm))
134 .collect::<Result<_, _>>()?,
135 origin,
136 ty,
137 span,
138 },
139
140 Let {
141 binding,
142 value,
143 mutable,
144 mut_span,
145 else_block,
146 else_span,
147 typed_pattern,
148 ty,
149 span,
150 } => Let {
151 binding,
152 value: Box::new(self.fold_expression(*value)?),
153 mutable,
154 mut_span,
155 else_block: else_block
156 .map(|e| self.fold_expression(*e).map(Box::new))
157 .transpose()?,
158 else_span,
159 typed_pattern,
160 ty,
161 span,
162 },
163
164 Return {
165 expression,
166 ty,
167 span,
168 } => Return {
169 expression: Box::new(self.fold_expression(*expression)?),
170 ty,
171 span,
172 },
173
174 Propagate {
175 expression,
176 ty,
177 span,
178 } => Propagate {
179 expression: Box::new(self.fold_expression(*expression)?),
180 ty,
181 span,
182 },
183
184 Unary {
185 operator,
186 expression,
187 ty,
188 span,
189 } => Unary {
190 operator,
191 expression: Box::new(self.fold_expression(*expression)?),
192 ty,
193 span,
194 },
195
196 Paren {
197 expression,
198 ty,
199 span,
200 } => Paren {
201 expression: Box::new(self.fold_expression(*expression)?),
202 ty,
203 span,
204 },
205
206 DotAccess {
207 expression,
208 member,
209 ty,
210 span,
211 dot_access_kind,
212 receiver_coercion,
213 } => DotAccess {
214 expression: Box::new(self.fold_expression(*expression)?),
215 member,
216 ty,
217 span,
218 dot_access_kind,
219 receiver_coercion,
220 },
221
222 IndexedAccess {
223 expression,
224 index,
225 ty,
226 span,
227 } => IndexedAccess {
228 expression: Box::new(self.fold_expression(*expression)?),
229 index: Box::new(self.fold_expression(*index)?),
230 ty,
231 span,
232 },
233
234 Assignment {
235 target,
236 value,
237 compound_operator,
238 span,
239 } => Assignment {
240 target: Box::new(self.fold_expression(*target)?),
241 value: Box::new(self.fold_expression(*value)?),
242 compound_operator,
243 span,
244 },
245
246 Tuple { elements, ty, span } => Tuple {
247 elements: self.fold_vec(elements)?,
248 ty,
249 span,
250 },
251
252 StructCall {
253 name,
254 field_assignments,
255 spread,
256 ty,
257 span,
258 } => StructCall {
259 name,
260 field_assignments: field_assignments
261 .into_iter()
262 .map(|mut f| {
263 f.value = Box::new(self.fold_expression(*f.value)?);
264 Ok(f)
265 })
266 .collect::<Result<_, Self::Error>>()?,
267 spread: Box::new((*spread).map(|e| self.fold_expression(e)).transpose()?),
268 ty,
269 span,
270 },
271
272 Function {
273 doc,
274 attributes,
275 name,
276 name_span,
277 generics,
278 params,
279 return_annotation,
280 return_type,
281 visibility,
282 body,
283 ty,
284 span,
285 } => Function {
286 doc,
287 attributes,
288 name,
289 name_span,
290 generics,
291 params,
292 return_annotation,
293 return_type,
294 visibility,
295 body: Box::new(self.fold_expression(*body)?),
296 ty,
297 span,
298 },
299
300 Lambda {
301 params,
302 return_annotation,
303 body,
304 ty,
305 span,
306 } => Lambda {
307 params,
308 return_annotation,
309 body: Box::new(self.fold_expression(*body)?),
310 ty,
311 span,
312 },
313
314 Reference {
315 expression,
316 ty,
317 span,
318 } => Reference {
319 expression: Box::new(self.fold_expression(*expression)?),
320 ty,
321 span,
322 },
323
324 For {
325 binding,
326 iterable,
327 body,
328 span,
329 needs_label,
330 } => For {
331 binding,
332 iterable: Box::new(self.fold_expression(*iterable)?),
333 body: Box::new(self.fold_expression(*body)?),
334 span,
335 needs_label,
336 },
337
338 While {
339 condition,
340 body,
341 span,
342 needs_label,
343 } => While {
344 condition: Box::new(self.fold_expression(*condition)?),
345 body: Box::new(self.fold_expression(*body)?),
346 span,
347 needs_label,
348 },
349
350 WhileLet {
351 pattern,
352 scrutinee,
353 body,
354 typed_pattern,
355 span,
356 needs_label,
357 } => WhileLet {
358 pattern,
359 scrutinee: Box::new(self.fold_expression(*scrutinee)?),
360 body: Box::new(self.fold_expression(*body)?),
361 typed_pattern,
362 span,
363 needs_label,
364 },
365
366 Loop {
367 body,
368 ty,
369 span,
370 needs_label,
371 } => Loop {
372 body: Box::new(self.fold_expression(*body)?),
373 ty,
374 span,
375 needs_label,
376 },
377
378 Task {
379 expression,
380 ty,
381 span,
382 } => Task {
383 expression: Box::new(self.fold_expression(*expression)?),
384 ty,
385 span,
386 },
387
388 Defer {
389 expression,
390 ty,
391 span,
392 } => Defer {
393 expression: Box::new(self.fold_expression(*expression)?),
394 ty,
395 span,
396 },
397
398 Select { arms, ty, span } => Select {
399 arms: arms
400 .into_iter()
401 .map(|arm| self.fold_select_arm(arm))
402 .collect::<Result<_, _>>()?,
403 ty,
404 span,
405 },
406
407 ImplBlock {
408 annotation,
409 receiver_name,
410 methods,
411 generics,
412 ty,
413 span,
414 } => ImplBlock {
415 annotation,
416 receiver_name,
417 methods: self.fold_vec(methods)?,
418 generics,
419 ty,
420 span,
421 },
422
423 Const {
424 doc,
425 identifier,
426 identifier_span,
427 annotation,
428 expression,
429 visibility,
430 ty,
431 span,
432 } => Const {
433 doc,
434 identifier,
435 identifier_span,
436 annotation,
437 expression: Box::new(self.fold_expression(*expression)?),
438 visibility,
439 ty,
440 span,
441 },
442
443 Cast {
444 expression,
445 target_type,
446 ty,
447 span,
448 } => Cast {
449 expression: Box::new(self.fold_expression(*expression)?),
450 target_type,
451 ty,
452 span,
453 },
454
455 Break { value, span } => Break {
456 value: match value {
457 Some(v) => Some(Box::new(self.fold_expression(*v)?)),
458 None => None,
459 },
460 span,
461 },
462
463 Literal {
464 literal: crate::ast::Literal::FormatString(parts),
465 ty,
466 span,
467 } => {
468 let folded_parts = parts
469 .into_iter()
470 .map(|part| match part {
471 FormatStringPart::Expression(expression) => {
472 Ok(FormatStringPart::Expression(Box::new(
473 self.fold_expression(*expression)?,
474 )))
475 }
476 other => Ok(other),
477 })
478 .collect::<Result<Vec<_>, Self::Error>>()?;
479 Literal {
480 literal: crate::ast::Literal::FormatString(folded_parts),
481 ty,
482 span,
483 }
484 }
485
486 Literal {
487 literal: crate::ast::Literal::Slice(elements),
488 ty,
489 span,
490 } => {
491 let folded_elements = self.fold_vec(elements)?;
492 Literal {
493 literal: crate::ast::Literal::Slice(folded_elements),
494 ty,
495 span,
496 }
497 }
498
499 Range {
500 start,
501 end,
502 inclusive,
503 ty,
504 span,
505 } => Range {
506 start: start
507 .map(|e| self.fold_expression(*e).map(Box::new))
508 .transpose()?,
509 end: end
510 .map(|e| self.fold_expression(*e).map(Box::new))
511 .transpose()?,
512 inclusive,
513 ty,
514 span,
515 },
516
517 Literal { .. }
518 | Identifier { .. }
519 | Enum { .. }
520 | ValueEnum { .. }
521 | Struct { .. }
522 | TypeAlias { .. }
523 | VariableDeclaration { .. }
524 | ModuleImport { .. }
525 | Interface { .. }
526 | Continue { .. }
527 | Unit { .. }
528 | RawGo { .. }
529 | NoOp => expression,
530 })
531 }
532
533 fn fold_vec(&mut self, expressions: Vec<Expression>) -> Result<Vec<Expression>, Self::Error> {
534 expressions
535 .into_iter()
536 .map(|e| self.fold_expression(e))
537 .collect()
538 }
539
540 fn fold_match_arm(&mut self, mut arm: MatchArm) -> Result<MatchArm, Self::Error> {
541 arm.expression = Box::new(self.fold_expression(*arm.expression)?);
542 arm.guard = arm
543 .guard
544 .map(|g| self.fold_expression(*g).map(Box::new))
545 .transpose()?;
546 Ok(arm)
547 }
548
549 fn fold_select_arm(&mut self, arm: SelectArm) -> Result<SelectArm, Self::Error> {
550 let pattern = match arm.pattern {
551 SelectArmPattern::Receive {
552 binding,
553 typed_pattern,
554 receive_expression,
555 body,
556 } => SelectArmPattern::Receive {
557 binding,
558 typed_pattern,
559 receive_expression: Box::new(self.fold_expression(*receive_expression)?),
560 body: Box::new(self.fold_expression(*body)?),
561 },
562 SelectArmPattern::Send {
563 send_expression,
564 body,
565 } => SelectArmPattern::Send {
566 send_expression: Box::new(self.fold_expression(*send_expression)?),
567 body: Box::new(self.fold_expression(*body)?),
568 },
569 SelectArmPattern::MatchReceive {
570 receive_expression,
571 arms,
572 } => SelectArmPattern::MatchReceive {
573 receive_expression: Box::new(self.fold_expression(*receive_expression)?),
574 arms: arms
575 .into_iter()
576 .map(|arm| self.fold_match_arm(arm))
577 .collect::<Result<_, _>>()?,
578 },
579 SelectArmPattern::WildCard { body } => SelectArmPattern::WildCard {
580 body: Box::new(self.fold_expression(*body)?),
581 },
582 };
583 Ok(SelectArm { pattern })
584 }
585}