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