1use crate::arena::Arena;
16use crate::ast::stmt::{
17 CaseTable, ClosureBody, DomQuery, Expr, MatchArm, ReadSource, SecurePad, SelectBranch, Stmt,
18 StringPart, TypeExpr,
19};
20
21#[derive(Clone, Copy)]
23pub struct Arenas<'a> {
24 pub exprs: &'a Arena<Expr<'a>>,
25 pub stmts: &'a Arena<Stmt<'a>>,
26 pub types: &'a Arena<TypeExpr<'a>>,
27}
28
29pub trait AstMap<'a> {
32 fn arenas(&self) -> Arenas<'a>;
33
34 fn map_expr(&mut self, e: &'a Expr<'a>) -> &'a Expr<'a> {
35 walk_expr(self, e)
36 }
37
38 fn map_stmt(&mut self, s: &Stmt<'a>) -> Stmt<'a> {
39 walk_stmt(self, s)
40 }
41
42 fn map_type(&mut self, t: &'a TypeExpr<'a>) -> &'a TypeExpr<'a> {
43 walk_type(self, t)
44 }
45
46 fn map_arm(&mut self, arm: &MatchArm<'a>) -> MatchArm<'a> {
47 walk_arm(self, arm)
48 }
49
50 fn map_block(&mut self, block: &'a [Stmt<'a>]) -> &'a [Stmt<'a>] {
53 walk_block(self, block)
54 }
55}
56
57pub fn walk_block<'a, M: AstMap<'a> + ?Sized>(m: &mut M, block: &'a [Stmt<'a>]) -> &'a [Stmt<'a>] {
58 let mapped: Vec<Stmt<'a>> = block.iter().map(|s| m.map_stmt(s)).collect();
59 m.arenas().stmts.alloc_slice(mapped)
60}
61
62fn map_exprs<'a, M: AstMap<'a> + ?Sized>(m: &mut M, exprs: &[&'a Expr<'a>]) -> Vec<&'a Expr<'a>> {
63 exprs.iter().map(|e| m.map_expr(e)).collect()
64}
65
66fn map_fields<'a, M: AstMap<'a> + ?Sized>(
67 m: &mut M,
68 fields: &[(crate::intern::Symbol, &'a Expr<'a>)],
69) -> Vec<(crate::intern::Symbol, &'a Expr<'a>)> {
70 fields.iter().map(|(f, e)| (*f, m.map_expr(e))).collect()
71}
72
73fn map_typed_params<'a, M: AstMap<'a> + ?Sized>(
74 m: &mut M,
75 params: &[(crate::intern::Symbol, &'a TypeExpr<'a>)],
76) -> Vec<(crate::intern::Symbol, &'a TypeExpr<'a>)> {
77 params.iter().map(|(p, t)| (*p, m.map_type(t))).collect()
78}
79
80fn map_type_slice<'a, M: AstMap<'a> + ?Sized>(
81 m: &mut M,
82 types: &'a [TypeExpr<'a>],
83) -> &'a [TypeExpr<'a>] {
84 let mapped: Vec<TypeExpr<'a>> = types.iter().map(|t| m.map_type(t).clone()).collect();
85 m.arenas().types.alloc_slice(mapped)
86}
87
88fn map_secure<'a, M: AstMap<'a> + ?Sized>(
89 m: &mut M,
90 pad: &Option<SecurePad<'a>>,
91) -> Option<SecurePad<'a>> {
92 pad.map(|p| SecurePad {
93 pad: m.map_expr(p.pad),
94 role: p.role,
95 })
96}
97
98pub fn walk_arm<'a, M: AstMap<'a> + ?Sized>(m: &mut M, arm: &MatchArm<'a>) -> MatchArm<'a> {
99 MatchArm {
100 enum_name: arm.enum_name,
101 variant: arm.variant,
102 bindings: arm.bindings.clone(),
103 body: m.map_block(arm.body),
104 }
105}
106
107pub fn walk_type<'a, M: AstMap<'a> + ?Sized>(m: &mut M, t: &'a TypeExpr<'a>) -> &'a TypeExpr<'a> {
108 let rebuilt = match t {
109 TypeExpr::Primitive(_) | TypeExpr::Named(_) => return t,
110 TypeExpr::Generic { base, params } => TypeExpr::Generic {
111 base: *base,
112 params: map_type_slice(m, params),
113 },
114 TypeExpr::Tuple(items) => TypeExpr::Tuple(map_type_slice(m, items)),
115 TypeExpr::Record {
116 field_names,
117 field_types,
118 } => TypeExpr::Record {
119 field_names,
120 field_types: map_type_slice(m, field_types),
121 },
122 TypeExpr::Function { inputs, output } => TypeExpr::Function {
123 inputs: map_type_slice(m, inputs),
124 output: m.map_type(output),
125 },
126 TypeExpr::Refinement {
127 base,
128 var,
129 predicate,
130 } => TypeExpr::Refinement {
131 base: m.map_type(base),
132 var: *var,
133 predicate,
134 },
135 TypeExpr::Persistent { inner } => TypeExpr::Persistent {
136 inner: m.map_type(inner),
137 },
138 TypeExpr::Mutable { inner } => TypeExpr::Mutable {
139 inner: m.map_type(inner),
140 },
141 };
142 m.arenas().types.alloc(rebuilt)
143}
144
145pub fn walk_expr<'a, M: AstMap<'a> + ?Sized>(m: &mut M, e: &'a Expr<'a>) -> &'a Expr<'a> {
146 let rebuilt = match e {
147 Expr::Literal(_) | Expr::Identifier(_) | Expr::Escape { .. } | Expr::OptionNone => {
148 return e
149 }
150 Expr::View(_) => return e,
152 Expr::BinaryOp { op, left, right } => Expr::BinaryOp {
153 op: *op,
154 left: m.map_expr(left),
155 right: m.map_expr(right),
156 },
157 Expr::Not { operand } => Expr::Not {
158 operand: m.map_expr(operand),
159 },
160 Expr::Call { function, args } => Expr::Call {
161 function: *function,
162 args: map_exprs(m, args),
163 },
164 Expr::Index { collection, index } => Expr::Index {
165 collection: m.map_expr(collection),
166 index: m.map_expr(index),
167 },
168 Expr::Slice {
169 collection,
170 start,
171 end,
172 } => Expr::Slice {
173 collection: m.map_expr(collection),
174 start: m.map_expr(start),
175 end: m.map_expr(end),
176 },
177 Expr::Copy { expr } => Expr::Copy {
178 expr: m.map_expr(expr),
179 },
180 Expr::Give { value } => Expr::Give {
181 value: m.map_expr(value),
182 },
183 Expr::Length { collection } => Expr::Length {
184 collection: m.map_expr(collection),
185 },
186 Expr::Contains { collection, value } => Expr::Contains {
187 collection: m.map_expr(collection),
188 value: m.map_expr(value),
189 },
190 Expr::Union { left, right } => Expr::Union {
191 left: m.map_expr(left),
192 right: m.map_expr(right),
193 },
194 Expr::Intersection { left, right } => Expr::Intersection {
195 left: m.map_expr(left),
196 right: m.map_expr(right),
197 },
198 Expr::ManifestOf { zone } => Expr::ManifestOf {
199 zone: m.map_expr(zone),
200 },
201 Expr::ChunkAt { index, zone } => Expr::ChunkAt {
202 index: m.map_expr(index),
203 zone: m.map_expr(zone),
204 },
205 Expr::List(items) => Expr::List(map_exprs(m, items)),
206 Expr::Tuple(items) => Expr::Tuple(map_exprs(m, items)),
207 Expr::Range { start, end } => Expr::Range {
208 start: m.map_expr(start),
209 end: m.map_expr(end),
210 },
211 Expr::FieldAccess { object, field } => Expr::FieldAccess {
212 object: m.map_expr(object),
213 field: *field,
214 },
215 Expr::New {
216 type_name,
217 type_args,
218 init_fields,
219 } => Expr::New {
220 type_name: *type_name,
221 type_args: type_args.iter().map(|t| m.map_type(t).clone()).collect(),
222 init_fields: map_fields(m, init_fields),
223 },
224 Expr::NewVariant {
225 enum_name,
226 variant,
227 fields,
228 } => Expr::NewVariant {
229 enum_name: *enum_name,
230 variant: *variant,
231 fields: map_fields(m, fields),
232 },
233 Expr::OptionSome { value } => Expr::OptionSome {
234 value: m.map_expr(value),
235 },
236 Expr::WithCapacity { value, capacity } => Expr::WithCapacity {
237 value: m.map_expr(value),
238 capacity: m.map_expr(capacity),
239 },
240 Expr::Closure {
241 params,
242 body,
243 return_type,
244 } => Expr::Closure {
245 params: map_typed_params(m, params),
246 body: match body {
247 ClosureBody::Expression(x) => ClosureBody::Expression(m.map_expr(x)),
248 ClosureBody::Block(b) => ClosureBody::Block(m.map_block(b)),
249 },
250 return_type: return_type.map(|t| m.map_type(t)),
251 },
252 Expr::CallExpr { callee, args } => Expr::CallExpr {
253 callee: m.map_expr(callee),
254 args: map_exprs(m, args),
255 },
256 Expr::InterpolatedString(parts) => Expr::InterpolatedString(
257 parts
258 .iter()
259 .map(|p| match p {
260 StringPart::Literal(s) => StringPart::Literal(*s),
261 StringPart::Expr {
262 value,
263 format_spec,
264 debug,
265 } => StringPart::Expr {
266 value: m.map_expr(value),
267 format_spec: *format_spec,
268 debug: *debug,
269 },
270 })
271 .collect(),
272 ),
273 Expr::Dom(q) => Expr::Dom(match q {
274 DomQuery::Text { selector } => DomQuery::Text {
275 selector: m.map_expr(selector),
276 },
277 DomQuery::Count { selector } => DomQuery::Count {
278 selector: m.map_expr(selector),
279 },
280 DomQuery::Value { selector } => DomQuery::Value {
281 selector: m.map_expr(selector),
282 },
283 DomQuery::Visible { selector } => DomQuery::Visible {
284 selector: m.map_expr(selector),
285 },
286 DomQuery::Hidden { selector } => DomQuery::Hidden {
287 selector: m.map_expr(selector),
288 },
289 DomQuery::Disabled { selector } => DomQuery::Disabled {
290 selector: m.map_expr(selector),
291 },
292 DomQuery::Checked { selector } => DomQuery::Checked {
293 selector: m.map_expr(selector),
294 },
295 DomQuery::Focused { selector } => DomQuery::Focused {
296 selector: m.map_expr(selector),
297 },
298 DomQuery::Journal => DomQuery::Journal,
299 DomQuery::JournalShows { first, second } => DomQuery::JournalShows {
300 first: *first,
301 second: *second,
302 },
303 DomQuery::Title => DomQuery::Title,
304 DomQuery::Url => DomQuery::Url,
305 }),
306 };
307 m.arenas().exprs.alloc(rebuilt)
308}
309
310pub fn walk_stmt<'a, M: AstMap<'a> + ?Sized>(m: &mut M, s: &Stmt<'a>) -> Stmt<'a> {
311 match s {
312 Stmt::Let {
313 var,
314 ty,
315 value,
316 mutable,
317 } => Stmt::Let {
318 var: *var,
319 ty: ty.map(|t| m.map_type(t)),
320 value: m.map_expr(value),
321 mutable: *mutable,
322 },
323 Stmt::Set { target, value } => Stmt::Set {
324 target: *target,
325 value: m.map_expr(value),
326 },
327 Stmt::Call { function, args } => Stmt::Call {
328 function: *function,
329 args: map_exprs(m, args),
330 },
331 Stmt::If {
332 cond,
333 then_block,
334 else_block,
335 } => Stmt::If {
336 cond: m.map_expr(cond),
337 then_block: m.map_block(then_block),
338 else_block: else_block.map(|b| m.map_block(b)),
339 },
340 Stmt::While {
341 cond,
342 body,
343 decreasing,
344 } => Stmt::While {
345 cond: m.map_expr(cond),
346 body: m.map_block(body),
347 decreasing: decreasing.map(|d| m.map_expr(d)),
348 },
349 Stmt::Repeat {
350 pattern,
351 iterable,
352 body,
353 } => Stmt::Repeat {
354 pattern: pattern.clone(),
355 iterable: m.map_expr(iterable),
356 body: m.map_block(body),
357 },
358 Stmt::Return { value } => Stmt::Return {
359 value: value.map(|v| m.map_expr(v)),
360 },
361 Stmt::RuntimeAssert {
362 condition,
363 hard,
364 dim,
365 reason,
366 } => Stmt::RuntimeAssert {
367 condition: m.map_expr(condition),
368 hard: *hard,
369 dim: *dim,
370 reason: reason.map(|r| m.map_expr(r)),
371 },
372 Stmt::Fail { reason } => Stmt::Fail {
373 reason: m.map_expr(reason),
374 },
375 Stmt::Give { object, recipient } => Stmt::Give {
376 object: m.map_expr(object),
377 recipient: m.map_expr(recipient),
378 },
379 Stmt::Show { object, recipient } => Stmt::Show {
380 object: m.map_expr(object),
381 recipient: m.map_expr(recipient),
382 },
383 Stmt::SetField {
384 object,
385 field,
386 value,
387 } => Stmt::SetField {
388 object: m.map_expr(object),
389 field: *field,
390 value: m.map_expr(value),
391 },
392 Stmt::FunctionDef {
393 name,
394 generics,
395 params,
396 body,
397 return_type,
398 is_native,
399 native_path,
400 is_exported,
401 export_target,
402 opt_flags,
403 } => Stmt::FunctionDef {
404 name: *name,
405 generics: generics.clone(),
406 params: map_typed_params(m, params),
407 body: m.map_block(body),
408 return_type: return_type.map(|t| m.map_type(t)),
409 is_native: *is_native,
410 native_path: *native_path,
411 is_exported: *is_exported,
412 export_target: *export_target,
413 opt_flags: *opt_flags,
414 },
415 Stmt::TestDef {
416 name,
417 skip,
418 capabilities,
419 budget,
420 cases,
421 body,
422 } => Stmt::TestDef {
423 name: *name,
424 skip: *skip,
425 capabilities: *capabilities,
426 budget: *budget,
427 cases: cases.as_ref().map(|c| CaseTable {
428 columns: c.columns.clone(),
429 rows: c.rows.iter().map(|row| map_exprs(m, row)).collect(),
430 }),
431 body: m.map_block(body),
432 },
433 Stmt::Expect { condition } => Stmt::Expect {
434 condition: m.map_expr(condition),
435 },
436 Stmt::ExpectFail { message, body } => Stmt::ExpectFail {
437 message: *message,
438 body: m.map_block(body),
439 },
440 Stmt::AdvanceTime { millis } => Stmt::AdvanceTime {
441 millis: m.map_expr(millis),
442 },
443 Stmt::Observe {
444 category,
445 message,
446 payloads,
447 sample,
448 } => Stmt::Observe {
449 category: *category,
450 message: *message,
451 payloads: map_fields(m, payloads),
452 sample: *sample,
453 },
454 Stmt::During { label, body } => Stmt::During {
455 label: *label,
456 body: m.map_block(body),
457 },
458 Stmt::ForAll {
459 vars,
460 domain,
461 constraint,
462 body,
463 } => Stmt::ForAll {
464 vars: vars.clone(),
465 domain: domain.clone(),
466 constraint: constraint.map(|c| m.map_expr(c)),
467 body: m.map_block(body),
468 },
469 Stmt::Inspect {
470 target,
471 arms,
472 has_otherwise,
473 } => Stmt::Inspect {
474 target: m.map_expr(target),
475 arms: arms.iter().map(|a| m.map_arm(a)).collect(),
476 has_otherwise: *has_otherwise,
477 },
478 Stmt::Push { value, collection } => Stmt::Push {
479 value: m.map_expr(value),
480 collection: m.map_expr(collection),
481 },
482 Stmt::Pop { collection, into } => Stmt::Pop {
483 collection: m.map_expr(collection),
484 into: *into,
485 },
486 Stmt::Add { value, collection } => Stmt::Add {
487 value: m.map_expr(value),
488 collection: m.map_expr(collection),
489 },
490 Stmt::Remove { value, collection } => Stmt::Remove {
491 value: m.map_expr(value),
492 collection: m.map_expr(collection),
493 },
494 Stmt::SetIndex {
495 collection,
496 index,
497 value,
498 } => Stmt::SetIndex {
499 collection: m.map_expr(collection),
500 index: m.map_expr(index),
501 value: m.map_expr(value),
502 },
503 Stmt::Splice { body } => Stmt::Splice {
504 body: m.map_block(body),
505 },
506 Stmt::Reversibly { body } => Stmt::Reversibly {
507 body: m.map_block(body),
508 },
509 Stmt::Zone {
510 name,
511 capacity,
512 source_file,
513 body,
514 } => Stmt::Zone {
515 name: *name,
516 capacity: *capacity,
517 source_file: source_file.clone(),
518 body: m.map_block(body),
519 },
520 Stmt::Concurrent { tasks } => Stmt::Concurrent {
521 tasks: m.map_block(tasks),
522 },
523 Stmt::Parallel { tasks } => Stmt::Parallel {
524 tasks: m.map_block(tasks),
525 },
526 Stmt::ReadFrom { var, source } => Stmt::ReadFrom {
527 var: *var,
528 source: match source {
529 ReadSource::Console => ReadSource::Console,
530 ReadSource::File(path) => ReadSource::File(m.map_expr(path)),
531 },
532 },
533 Stmt::WriteFile { content, path } => Stmt::WriteFile {
534 content: m.map_expr(content),
535 path: m.map_expr(path),
536 },
537 Stmt::SendMessage {
538 message,
539 destination,
540 compression,
541 cached,
542 unchecked,
543 layout,
544 shared,
545 computed,
546 indexed,
547 deduped,
548 unreliable,
549 } => Stmt::SendMessage {
550 message: m.map_expr(message),
551 destination: m.map_expr(destination),
552 compression: *compression,
553 cached: *cached,
554 unchecked: *unchecked,
555 layout: *layout,
556 shared: *shared,
557 computed: *computed,
558 indexed: *indexed,
559 deduped: *deduped,
560 unreliable: *unreliable,
561 },
562 Stmt::AwaitMessage {
563 source,
564 into,
565 view,
566 stream,
567 } => Stmt::AwaitMessage {
568 source: m.map_expr(source),
569 into: *into,
570 view: *view,
571 stream: *stream,
572 },
573 Stmt::StreamMessage {
574 values,
575 destination,
576 } => Stmt::StreamMessage {
577 values: m.map_expr(values),
578 destination: m.map_expr(destination),
579 },
580 Stmt::MergeCrdt { source, target } => Stmt::MergeCrdt {
581 source: m.map_expr(source),
582 target: m.map_expr(target),
583 },
584 Stmt::JoinCrdt { source, target } => Stmt::JoinCrdt {
585 source: m.map_expr(source),
586 target: m.map_expr(target),
587 },
588 Stmt::IncreaseCrdt {
589 object,
590 field,
591 amount,
592 } => Stmt::IncreaseCrdt {
593 object: m.map_expr(object),
594 field: *field,
595 amount: m.map_expr(amount),
596 },
597 Stmt::DecreaseCrdt {
598 object,
599 field,
600 amount,
601 } => Stmt::DecreaseCrdt {
602 object: m.map_expr(object),
603 field: *field,
604 amount: m.map_expr(amount),
605 },
606 Stmt::AppendToSequence { sequence, value } => Stmt::AppendToSequence {
607 sequence: m.map_expr(sequence),
608 value: m.map_expr(value),
609 },
610 Stmt::ResolveConflict {
611 object,
612 field,
613 value,
614 } => Stmt::ResolveConflict {
615 object: m.map_expr(object),
616 field: *field,
617 value: m.map_expr(value),
618 },
619 Stmt::Listen { address, secure } => Stmt::Listen {
620 address: m.map_expr(address),
621 secure: map_secure(m, secure),
622 },
623 Stmt::Serve { port, handler } => Stmt::Serve {
624 port: m.map_expr(port),
625 handler: m.map_expr(handler),
626 },
627 Stmt::ConnectTo { address, secure } => Stmt::ConnectTo {
628 address: m.map_expr(address),
629 secure: map_secure(m, secure),
630 },
631 Stmt::LetPeerAgent { var, address } => Stmt::LetPeerAgent {
632 var: *var,
633 address: m.map_expr(address),
634 },
635 Stmt::Sleep { milliseconds } => Stmt::Sleep {
636 milliseconds: m.map_expr(milliseconds),
637 },
638 Stmt::Sync { var, topic } => Stmt::Sync {
639 var: *var,
640 topic: m.map_expr(topic),
641 },
642 Stmt::Fetch {
643 method,
644 url,
645 body,
646 as_type,
647 then,
648 } => Stmt::Fetch {
649 method: *method,
650 url: m.map_expr(url),
651 body: body.map(|b| m.map_expr(b)),
652 as_type: m.map_type(as_type),
653 then: *then,
654 },
655 Stmt::GivenFetch {
656 method,
657 url,
658 status,
659 body,
660 } => Stmt::GivenFetch {
661 method: *method,
662 url: m.map_expr(url),
663 status: m.map_expr(status),
664 body: m.map_expr(body),
665 },
666 Stmt::DirectiveDef {
667 name,
668 params,
669 body,
670 removed,
671 } => Stmt::DirectiveDef {
672 name: *name,
673 params: map_typed_params(m, params),
674 body: m.map_block(body),
675 removed: removed.map(|b| m.map_block(b)),
676 },
677 Stmt::MiddlewareDef { param, body } => Stmt::MiddlewareDef {
678 param: (param.0, m.map_type(param.1)),
679 body: m.map_block(body),
680 },
681 Stmt::Mount { var, path } => Stmt::Mount {
682 var: *var,
683 path: m.map_expr(path),
684 },
685 Stmt::LaunchTask { function, args } => Stmt::LaunchTask {
686 function: *function,
687 args: map_exprs(m, args),
688 },
689 Stmt::LaunchTaskWithHandle {
690 handle,
691 function,
692 args,
693 } => Stmt::LaunchTaskWithHandle {
694 handle: *handle,
695 function: *function,
696 args: map_exprs(m, args),
697 },
698 Stmt::SendPipe { value, pipe } => Stmt::SendPipe {
699 value: m.map_expr(value),
700 pipe: m.map_expr(pipe),
701 },
702 Stmt::ReceivePipe { var, pipe } => Stmt::ReceivePipe {
703 var: *var,
704 pipe: m.map_expr(pipe),
705 },
706 Stmt::TrySendPipe {
707 value,
708 pipe,
709 result,
710 } => Stmt::TrySendPipe {
711 value: m.map_expr(value),
712 pipe: m.map_expr(pipe),
713 result: *result,
714 },
715 Stmt::TryReceivePipe { var, pipe } => Stmt::TryReceivePipe {
716 var: *var,
717 pipe: m.map_expr(pipe),
718 },
719 Stmt::StopTask { handle } => Stmt::StopTask {
720 handle: m.map_expr(handle),
721 },
722 Stmt::CloseChannel { channel } => Stmt::CloseChannel {
723 channel: m.map_expr(channel),
724 },
725 Stmt::AwaitTask { handle, result } => Stmt::AwaitTask {
726 handle: m.map_expr(handle),
727 result: *result,
728 },
729 Stmt::Select { branches } => Stmt::Select {
730 branches: branches
731 .iter()
732 .map(|b| match b {
733 SelectBranch::Receive { var, pipe, body } => SelectBranch::Receive {
734 var: *var,
735 pipe: m.map_expr(pipe),
736 body: m.map_block(body),
737 },
738 SelectBranch::Timeout { milliseconds, body } => SelectBranch::Timeout {
739 milliseconds: m.map_expr(milliseconds),
740 body: m.map_block(body),
741 },
742 })
743 .collect(),
744 },
745 Stmt::Swap { .. }
747 | Stmt::Break
748 | Stmt::StructDef { .. }
749 | Stmt::ExpectOutput { .. }
750 | Stmt::GivenFile { .. }
751 | Stmt::Check { .. }
752 | Stmt::Spawn { .. }
753 | Stmt::Replay { .. }
754 | Stmt::ServeFetches
755 | Stmt::RedirectDef { .. }
756 | Stmt::CreatePipe { .. }
757 | Stmt::Escape { .. }
758 | Stmt::Require { .. }
759 | Stmt::Site(_) => s.clone(),
760 Stmt::Assert { .. }
762 | Stmt::Trust { .. }
763 | Stmt::Theorem(_)
764 | Stmt::Definition(_)
765 | Stmt::Axiom(_)
766 | Stmt::Theory(_) => s.clone(),
767 Stmt::Ui(_) | Stmt::ComponentDef(_) | Stmt::RouteDef { .. } => s.clone(),
769 }
770}