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 binding_id,
341 } => For {
342 binding,
343 iterable: Box::new(self.fold_expression(*iterable)?),
344 body: Box::new(self.fold_expression(*body)?),
345 span,
346 needs_label,
347 binding_id,
348 },
349
350 While {
351 condition,
352 body,
353 span,
354 needs_label,
355 } => While {
356 condition: Box::new(self.fold_expression(*condition)?),
357 body: Box::new(self.fold_expression(*body)?),
358 span,
359 needs_label,
360 },
361
362 WhileLet {
363 pattern,
364 scrutinee,
365 body,
366 typed_pattern,
367 span,
368 needs_label,
369 } => WhileLet {
370 pattern,
371 scrutinee: Box::new(self.fold_expression(*scrutinee)?),
372 body: Box::new(self.fold_expression(*body)?),
373 typed_pattern,
374 span,
375 needs_label,
376 },
377
378 Loop {
379 body,
380 ty,
381 span,
382 needs_label,
383 } => Loop {
384 body: Box::new(self.fold_expression(*body)?),
385 ty,
386 span,
387 needs_label,
388 },
389
390 Task {
391 expression,
392 ty,
393 span,
394 } => Task {
395 expression: Box::new(self.fold_expression(*expression)?),
396 ty,
397 span,
398 },
399
400 Defer {
401 expression,
402 ty,
403 span,
404 } => Defer {
405 expression: Box::new(self.fold_expression(*expression)?),
406 ty,
407 span,
408 },
409
410 Select { arms, ty, span } => Select {
411 arms: arms
412 .into_iter()
413 .map(|arm| self.fold_select_arm(arm))
414 .collect::<Result<_, _>>()?,
415 ty,
416 span,
417 },
418
419 ImplBlock {
420 annotation,
421 receiver_name,
422 methods,
423 generics,
424 ty,
425 span,
426 } => ImplBlock {
427 annotation,
428 receiver_name,
429 methods: self.fold_vec(methods)?,
430 generics,
431 ty,
432 span,
433 },
434
435 Const {
436 doc,
437 identifier,
438 identifier_span,
439 annotation,
440 expression,
441 visibility,
442 ty,
443 span,
444 } => Const {
445 doc,
446 identifier,
447 identifier_span,
448 annotation,
449 expression: Box::new(self.fold_expression(*expression)?),
450 visibility,
451 ty,
452 span,
453 },
454
455 Cast {
456 expression,
457 target_type,
458 ty,
459 span,
460 } => Cast {
461 expression: Box::new(self.fold_expression(*expression)?),
462 target_type,
463 ty,
464 span,
465 },
466
467 Break { value, span } => Break {
468 value: match value {
469 Some(v) => Some(Box::new(self.fold_expression(*v)?)),
470 None => None,
471 },
472 span,
473 },
474
475 Literal {
476 literal: crate::ast::Literal::FormatString(parts),
477 ty,
478 span,
479 } => {
480 let folded_parts = parts
481 .into_iter()
482 .map(|part| match part {
483 FormatStringPart::Expression(expression) => {
484 Ok(FormatStringPart::Expression(Box::new(
485 self.fold_expression(*expression)?,
486 )))
487 }
488 other => Ok(other),
489 })
490 .collect::<Result<Vec<_>, Self::Error>>()?;
491 Literal {
492 literal: crate::ast::Literal::FormatString(folded_parts),
493 ty,
494 span,
495 }
496 }
497
498 Literal {
499 literal: crate::ast::Literal::Slice(elements),
500 ty,
501 span,
502 } => {
503 let folded_elements = self.fold_vec(elements)?;
504 Literal {
505 literal: crate::ast::Literal::Slice(folded_elements),
506 ty,
507 span,
508 }
509 }
510
511 Range {
512 start,
513 end,
514 inclusive,
515 ty,
516 span,
517 } => Range {
518 start: start
519 .map(|e| self.fold_expression(*e).map(Box::new))
520 .transpose()?,
521 end: end
522 .map(|e| self.fold_expression(*e).map(Box::new))
523 .transpose()?,
524 inclusive,
525 ty,
526 span,
527 },
528
529 Literal { .. }
530 | Identifier { .. }
531 | Enum { .. }
532 | ValueEnum { .. }
533 | Struct { .. }
534 | TypeAlias { .. }
535 | VariableDeclaration { .. }
536 | ModuleImport { .. }
537 | Interface { .. }
538 | Continue { .. }
539 | Unit { .. }
540 | RawGo { .. }
541 | NoOp => expression,
542 })
543 }
544
545 fn fold_vec(&mut self, expressions: Vec<Expression>) -> Result<Vec<Expression>, Self::Error> {
546 expressions
547 .into_iter()
548 .map(|e| self.fold_expression(e))
549 .collect()
550 }
551
552 fn fold_match_arm(&mut self, mut arm: MatchArm) -> Result<MatchArm, Self::Error> {
553 arm.expression = Box::new(self.fold_expression(*arm.expression)?);
554 arm.guard = arm
555 .guard
556 .map(|g| self.fold_expression(*g).map(Box::new))
557 .transpose()?;
558 Ok(arm)
559 }
560
561 fn fold_select_arm(&mut self, arm: SelectArm) -> Result<SelectArm, Self::Error> {
562 let pattern = match arm.pattern {
563 SelectArmPattern::Receive {
564 binding,
565 typed_pattern,
566 receive_expression,
567 body,
568 } => SelectArmPattern::Receive {
569 binding,
570 typed_pattern,
571 receive_expression: Box::new(self.fold_expression(*receive_expression)?),
572 body: Box::new(self.fold_expression(*body)?),
573 },
574 SelectArmPattern::Send {
575 send_expression,
576 body,
577 } => SelectArmPattern::Send {
578 send_expression: Box::new(self.fold_expression(*send_expression)?),
579 body: Box::new(self.fold_expression(*body)?),
580 },
581 SelectArmPattern::MatchReceive {
582 receive_expression,
583 arms,
584 } => SelectArmPattern::MatchReceive {
585 receive_expression: Box::new(self.fold_expression(*receive_expression)?),
586 arms: arms
587 .into_iter()
588 .map(|arm| self.fold_match_arm(arm))
589 .collect::<Result<_, _>>()?,
590 },
591 SelectArmPattern::WildCard { body } => SelectArmPattern::WildCard {
592 body: Box::new(self.fold_expression(*body)?),
593 },
594 };
595 Ok(SelectArm { pattern })
596 }
597}