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