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 } => IndexedAccess {
230 expression: Box::new(self.fold_expression(*expression)?),
231 index: Box::new(self.fold_expression(*index)?),
232 ty,
233 span,
234 },
235
236 Assignment {
237 target,
238 value,
239 compound_operator,
240 span,
241 } => Assignment {
242 target: Box::new(self.fold_expression(*target)?),
243 value: Box::new(self.fold_expression(*value)?),
244 compound_operator,
245 span,
246 },
247
248 Tuple { elements, ty, span } => Tuple {
249 elements: self.fold_vec(elements)?,
250 ty,
251 span,
252 },
253
254 StructCall {
255 name,
256 field_assignments,
257 spread,
258 ty,
259 span,
260 } => StructCall {
261 name,
262 field_assignments: field_assignments
263 .into_iter()
264 .map(|mut f| {
265 f.value = Box::new(self.fold_expression(*f.value)?);
266 Ok(f)
267 })
268 .collect::<Result<_, Self::Error>>()?,
269 spread: match spread {
270 StructSpread::None => StructSpread::None,
271 StructSpread::From(e) => {
272 StructSpread::From(Box::new(self.fold_expression(*e)?))
273 }
274 StructSpread::ZeroFill { span } => StructSpread::ZeroFill { span },
275 },
276 ty,
277 span,
278 },
279
280 Function {
281 doc,
282 attributes,
283 name,
284 name_span,
285 generics,
286 params,
287 return_annotation,
288 return_type,
289 visibility,
290 body,
291 ty,
292 span,
293 } => Function {
294 doc,
295 attributes,
296 name,
297 name_span,
298 generics,
299 params,
300 return_annotation,
301 return_type,
302 visibility,
303 body: Box::new(self.fold_expression(*body)?),
304 ty,
305 span,
306 },
307
308 Lambda {
309 params,
310 return_annotation,
311 body,
312 ty,
313 span,
314 } => Lambda {
315 params,
316 return_annotation,
317 body: Box::new(self.fold_expression(*body)?),
318 ty,
319 span,
320 },
321
322 Reference {
323 expression,
324 ty,
325 span,
326 } => Reference {
327 expression: Box::new(self.fold_expression(*expression)?),
328 ty,
329 span,
330 },
331
332 For {
333 binding,
334 iterable,
335 body,
336 span,
337 needs_label,
338 } => For {
339 binding,
340 iterable: Box::new(self.fold_expression(*iterable)?),
341 body: Box::new(self.fold_expression(*body)?),
342 span,
343 needs_label,
344 },
345
346 While {
347 condition,
348 body,
349 span,
350 needs_label,
351 } => While {
352 condition: Box::new(self.fold_expression(*condition)?),
353 body: Box::new(self.fold_expression(*body)?),
354 span,
355 needs_label,
356 },
357
358 WhileLet {
359 pattern,
360 scrutinee,
361 body,
362 typed_pattern,
363 span,
364 needs_label,
365 } => WhileLet {
366 pattern,
367 scrutinee: Box::new(self.fold_expression(*scrutinee)?),
368 body: Box::new(self.fold_expression(*body)?),
369 typed_pattern,
370 span,
371 needs_label,
372 },
373
374 Loop {
375 body,
376 ty,
377 span,
378 needs_label,
379 } => Loop {
380 body: Box::new(self.fold_expression(*body)?),
381 ty,
382 span,
383 needs_label,
384 },
385
386 Task {
387 expression,
388 ty,
389 span,
390 } => Task {
391 expression: Box::new(self.fold_expression(*expression)?),
392 ty,
393 span,
394 },
395
396 Defer {
397 expression,
398 ty,
399 span,
400 } => Defer {
401 expression: Box::new(self.fold_expression(*expression)?),
402 ty,
403 span,
404 },
405
406 Select { arms, ty, span } => Select {
407 arms: arms
408 .into_iter()
409 .map(|arm| self.fold_select_arm(arm))
410 .collect::<Result<_, _>>()?,
411 ty,
412 span,
413 },
414
415 ImplBlock {
416 annotation,
417 receiver_name,
418 methods,
419 generics,
420 ty,
421 span,
422 } => ImplBlock {
423 annotation,
424 receiver_name,
425 methods: self.fold_vec(methods)?,
426 generics,
427 ty,
428 span,
429 },
430
431 Const {
432 doc,
433 identifier,
434 identifier_span,
435 annotation,
436 expression,
437 visibility,
438 ty,
439 span,
440 } => Const {
441 doc,
442 identifier,
443 identifier_span,
444 annotation,
445 expression: Box::new(self.fold_expression(*expression)?),
446 visibility,
447 ty,
448 span,
449 },
450
451 Cast {
452 expression,
453 target_type,
454 ty,
455 span,
456 } => Cast {
457 expression: Box::new(self.fold_expression(*expression)?),
458 target_type,
459 ty,
460 span,
461 },
462
463 Break { value, span } => Break {
464 value: match value {
465 Some(v) => Some(Box::new(self.fold_expression(*v)?)),
466 None => None,
467 },
468 span,
469 },
470
471 Literal {
472 literal: crate::ast::Literal::FormatString(parts),
473 ty,
474 span,
475 } => {
476 let folded_parts = parts
477 .into_iter()
478 .map(|part| match part {
479 FormatStringPart::Expression(expression) => {
480 Ok(FormatStringPart::Expression(Box::new(
481 self.fold_expression(*expression)?,
482 )))
483 }
484 other => Ok(other),
485 })
486 .collect::<Result<Vec<_>, Self::Error>>()?;
487 Literal {
488 literal: crate::ast::Literal::FormatString(folded_parts),
489 ty,
490 span,
491 }
492 }
493
494 Literal {
495 literal: crate::ast::Literal::Slice(elements),
496 ty,
497 span,
498 } => {
499 let folded_elements = self.fold_vec(elements)?;
500 Literal {
501 literal: crate::ast::Literal::Slice(folded_elements),
502 ty,
503 span,
504 }
505 }
506
507 Range {
508 start,
509 end,
510 inclusive,
511 ty,
512 span,
513 } => Range {
514 start: start
515 .map(|e| self.fold_expression(*e).map(Box::new))
516 .transpose()?,
517 end: end
518 .map(|e| self.fold_expression(*e).map(Box::new))
519 .transpose()?,
520 inclusive,
521 ty,
522 span,
523 },
524
525 Literal { .. }
526 | Identifier { .. }
527 | Enum { .. }
528 | ValueEnum { .. }
529 | Struct { .. }
530 | TypeAlias { .. }
531 | VariableDeclaration { .. }
532 | ModuleImport { .. }
533 | Interface { .. }
534 | Continue { .. }
535 | Unit { .. }
536 | RawGo { .. }
537 | NoOp => expression,
538 })
539 }
540
541 fn fold_vec(&mut self, expressions: Vec<Expression>) -> Result<Vec<Expression>, Self::Error> {
542 expressions
543 .into_iter()
544 .map(|e| self.fold_expression(e))
545 .collect()
546 }
547
548 fn fold_match_arm(&mut self, mut arm: MatchArm) -> Result<MatchArm, Self::Error> {
549 arm.expression = Box::new(self.fold_expression(*arm.expression)?);
550 arm.guard = arm
551 .guard
552 .map(|g| self.fold_expression(*g).map(Box::new))
553 .transpose()?;
554 Ok(arm)
555 }
556
557 fn fold_select_arm(&mut self, arm: SelectArm) -> Result<SelectArm, Self::Error> {
558 let pattern = match arm.pattern {
559 SelectArmPattern::Receive {
560 binding,
561 typed_pattern,
562 receive_expression,
563 body,
564 } => SelectArmPattern::Receive {
565 binding,
566 typed_pattern,
567 receive_expression: Box::new(self.fold_expression(*receive_expression)?),
568 body: Box::new(self.fold_expression(*body)?),
569 },
570 SelectArmPattern::Send {
571 send_expression,
572 body,
573 } => SelectArmPattern::Send {
574 send_expression: Box::new(self.fold_expression(*send_expression)?),
575 body: Box::new(self.fold_expression(*body)?),
576 },
577 SelectArmPattern::MatchReceive {
578 receive_expression,
579 arms,
580 } => SelectArmPattern::MatchReceive {
581 receive_expression: Box::new(self.fold_expression(*receive_expression)?),
582 arms: arms
583 .into_iter()
584 .map(|arm| self.fold_match_arm(arm))
585 .collect::<Result<_, _>>()?,
586 },
587 SelectArmPattern::WildCard { body } => SelectArmPattern::WildCard {
588 body: Box::new(self.fold_expression(*body)?),
589 },
590 };
591 Ok(SelectArm { pattern })
592 }
593}