mz_expr/visit.rs
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Visitor support for recursive data types.
11//!
12//! Recursive types can implement the [`VisitChildren`] trait, to
13//! specify how their recursive entries can be accessed. The extension
14//! trait [`Visit`] then adds support for iteratively traversing
15//! instances of those types.
16//!
17//! # Naming
18//!
19//! Visitor methods follow this naming pattern:
20//!
21//! ```text
22//! [try_]visit_[mut_]{children,post,pre}
23//! ```
24//!
25//! * The `try`-prefix specifies whether the visitor callback is
26//! fallible (prefix present) or infallible (prefix omitted).
27//! * The `mut`-suffix specifies whether the visitor callback gets
28//! access to mutable (prefix present) or immutable (prefix omitted)
29//! child references.
30//! * The final suffix determines the nature of the traversal:
31//! * `children`: only visit direct children
32//! * `post`: recursively visit children in post-order
33//! * `pre`: recursively visit children in pre-order
34//! * no suffix: recursively visit children in pre- and post-order
35//! using a ~Visitor~` that encapsulates the shared context.
36
37/// A trait for types that can visit their direct children of type `T`.
38///
39/// Implementing [`VisitChildren<Self>`] automatically also implements
40/// the [`Visit`] trait, which enables recursive traversal.
41///
42/// Note that care needs to be taken when implementing this trait for
43/// mutually recursive types (such as a type A where A has children
44/// of type B and vice versa). More specifically, at the moment it is
45/// not possible to implement versions of `VisitChildren<A> for A` such
46/// that A considers as its children all A-nodes occurring at leaf
47/// positions of B-children and vice versa for `VisitChildren<B> for B`.
48/// Doing this will result in recursion limit violations as indicated
49/// in the accompanying `test_recursive_types_b` test.
50pub trait VisitChildren<T> {
51 /// Apply an infallible immutable function `f` to each direct child.
52 fn visit_children<F>(&self, f: F)
53 where
54 F: FnMut(&T),
55 {
56 self.children().for_each(f);
57 }
58
59 /// Apply an infallible mutable function `f` to each direct child.
60 fn visit_mut_children<F>(&mut self, f: F)
61 where
62 F: FnMut(&mut T),
63 {
64 self.children_mut().for_each(f);
65 }
66
67 /// Apply a fallible immutable function `f` to each direct child.
68 fn try_visit_children<F, E>(&self, mut f: F) -> Result<(), E>
69 where
70 F: FnMut(&T) -> Result<(), E>,
71 {
72 for child in self.children() {
73 f(child)?;
74 }
75
76 Ok(())
77 }
78
79 /// Apply a fallible mutable function `f` to each direct child.
80 fn try_visit_mut_children<F, E>(&mut self, mut f: F) -> Result<(), E>
81 where
82 F: FnMut(&mut T) -> Result<(), E>,
83 {
84 for child in self.children_mut() {
85 f(child)?;
86 }
87
88 Ok(())
89 }
90
91 /// The `T`-typed children of this element.
92 fn children<'a>(&'a self) -> impl DoubleEndedIterator<Item = &'a T>
93 where
94 T: 'a;
95
96 /// The `&mut T`-typed children of this element.
97 ///
98 /// It is critical for the safety of mutable post-order traversals that this
99 /// function be written using safe code.
100 fn children_mut<'a>(&'a mut self) -> impl DoubleEndedIterator<Item = &'a mut T>
101 where
102 T: 'a;
103}
104
105/// A trait for types that can recursively visit their children of the
106/// same type.
107///
108/// This trait is automatically implemented for all implementors of
109/// [`VisitChildren`].
110///
111/// All methods provided by this trait are iterative.
112///
113/// The immutable visitors hand the callback references borrowed for the whole
114/// lifetime of `&self`, so a callback may keep them: accumulate them, or embed
115/// them in its error type. This is what lets a bottom-up fold report an error
116/// that points at the offending node.
117///
118/// The mutable visitors deliberately do not do this. Their callbacks are
119/// higher-ranked over the reference lifetime, which prevents a callback from
120/// retaining what it is handed. For the post-order and pre-post visitors that
121/// is load-bearing: they rebuild `&mut Self` from raw pointers while a parent's
122/// pointer stays on the traversal stack, so retaining a child reference, were
123/// that permitted, would alias `&mut` references to a parent and its child.
124/// See `VisitMutAction` for the full argument. The pre-order visitors keep
125/// plain `&mut Self` on their stack and consume a parent before visiting its
126/// children, so for them the borrow checker would reject the generalization
127/// rather than accept unsound code.
128///
129/// NB that any visitor with mutable post-traversal uses unsafe code. It is critical
130/// that `VisitChildren::children_mut` be written using safe code, i.e., no aliasing
131/// of children or access to parents.
132pub trait Visit {
133 /// Post-order immutable infallible visitor for `self`.
134 fn visit_post<'a, F>(&'a self, f: &mut F)
135 where
136 F: FnMut(&'a Self);
137
138 /// Post-order mutable infallible visitor for `self`.
139 fn visit_mut_post<F>(&mut self, f: &mut F)
140 where
141 F: FnMut(&mut Self);
142
143 /// Post-order immutable fallible visitor for `self`.
144 fn try_visit_post<'a, F, E>(&'a self, f: &mut F) -> Result<(), E>
145 where
146 F: FnMut(&'a Self) -> Result<(), E>;
147
148 /// Post-order mutable fallible visitor for `self`.
149 fn try_visit_mut_post<F, E>(&mut self, f: &mut F) -> Result<(), E>
150 where
151 F: FnMut(&mut Self) -> Result<(), E>;
152
153 /// Pre-order immutable infallible visitor for `self`.
154 fn visit_pre<'a, F>(&'a self, f: &mut F)
155 where
156 F: FnMut(&'a Self);
157
158 /// Pre-order immutable infallible visitor for `self`, which also accumulates context
159 /// information along the path from the root to the current node's parent.
160 /// `acc_fun` is a similar closure as in `fold`. The accumulated context is passed to the
161 /// visitor, along with the current node.
162 ///
163 /// For example, one can use this on a `MirScalarExpr` to tell the visitor whether the current
164 /// subexpression has a negation somewhere above it.
165 ///
166 /// When using it on a `MirRelationExpr`, one has to be mindful that `Let` bindings are not
167 /// followed, i.e., the context won't include what happens with a `Let` binding in some other
168 /// `MirRelationExpr` where the binding occurs in a `Get`.
169 fn visit_pre_with_context<'a, Context, AccFun, Visitor>(
170 &'a self,
171 init: Context,
172 acc_fun: &mut AccFun,
173 visitor: &mut Visitor,
174 ) where
175 Context: Clone,
176 AccFun: FnMut(Context, &'a Self) -> Context,
177 Visitor: FnMut(&Context, &'a Self);
178
179 /// Pre-order mutable infallible visitor for `self`.
180 fn visit_mut_pre<F>(&mut self, f: &mut F)
181 where
182 F: FnMut(&mut Self);
183
184 /// Pre-order immutable fallible visitor for `self`.
185 fn try_visit_pre<'a, F, E>(&'a self, f: &mut F) -> Result<(), E>
186 where
187 F: FnMut(&'a Self) -> Result<(), E>;
188
189 /// Pre-order mutable fallible visitor for `self`.
190 fn try_visit_mut_pre<F, E>(&mut self, f: &mut F) -> Result<(), E>
191 where
192 F: FnMut(&mut Self) -> Result<(), E>;
193
194 /// A generalization of [`Visit::visit_pre`] and [`Visit::visit_post`].
195 ///
196 /// The function `pre` runs on `self` before it runs on any of the children.
197 /// The function `post` runs on children first before the parent.
198 ///
199 /// Optionally, `pre` can return which children, if any, should be visited
200 /// (default is to visit all children).
201 fn visit_pre_post<'a, F1, F2>(&'a self, pre: &mut F1, post: &mut F2)
202 where
203 F1: FnMut(&'a Self) -> Option<Vec<&'a Self>>,
204 F2: FnMut(&'a Self);
205
206 /// A generalization of [`Visit::visit_mut_pre`] and [`Visit::visit_mut_post`].
207 ///
208 /// The function `pre` runs on `self` before it runs on any of the children.
209 /// The function `post` runs on children first before the parent.
210 ///
211 /// Optionally, `pre` can return which children, if any, should be visited
212 /// (default is to visit all children).
213 ///
214 /// It is important for safety that `pre` is (a) safe code and (b) returns children only.
215 fn visit_mut_pre_post<F1, F2>(&mut self, pre: &mut F1, post: &mut F2)
216 where
217 F1: FnMut(&mut Self) -> Option<Vec<&mut Self>>,
218 F2: FnMut(&mut Self);
219}
220
221/// Frames for immutable post-traversals, will be kept in a stack.
222enum VisitAction<'a, T> {
223 /// Put on the stack when entering a node.
224 ///
225 /// Causes us to push children.
226 Enter(&'a T),
227 /// Put on the stack when leaving a node, all children visited.
228 ///
229 /// Causes us to do the post-traversal visit of the parent.
230 Leave(&'a T),
231}
232
233/// Frames for mutable post-traversals, will be kept in a stack.
234///
235/// Notice that we use mutable pointers, because mutable post-traversal is unsafe in rust.
236///
237/// The core argument for correctness mirrors the tree-borrow correctness argument Rust's
238/// borrow checker uses for the ordinary function call stack. Loosely:
239///
240/// - We split a mutable parent node into its children, and put them on the stack.
241/// - We keep a mutable pointer to the parent node, but won't touch it until we're done with all children.
242/// + In the function call stack, this mutable pointer is the stack frame, safely inaccessible.
243/// + In our `unsafe` action stack, this mutable pointer is below all of the `Enter` actions,
244/// and we promise not to touch it until they complete.
245/// - When we have processed all children, we can reassemble access to the parent from its parts.
246/// + In the function call stack, we do this on return from recursive calls.
247/// - In our `unsafe` action stack, we do this after popping all children.
248enum VisitMutAction<T> {
249 /// Put on the stack when entering a node.
250 ///
251 /// Causes us to push children.
252 Enter(*mut T),
253 /// Put on the stack when leaving a node, all children visited.
254 ///
255 /// Causes us to do the post-traversal visit of the parent.
256 Leave(*mut T),
257}
258
259impl<T: VisitChildren<T>> Visit for T {
260 fn visit_post<'a, F>(&'a self, f: &mut F)
261 where
262 F: FnMut(&'a Self),
263 {
264 use VisitAction::*;
265 let mut stack = vec![Enter(self)];
266 while let Some(action) = stack.pop() {
267 match action {
268 Enter(elt) => {
269 stack.push(Leave(elt));
270 // Push children in reverse so they pop (and are visited) left-to-right.
271 stack.extend(elt.children().rev().map(Enter));
272 }
273 Leave(elt) => f(elt),
274 }
275 }
276 }
277
278 #[allow(clippy::as_conversions)]
279 fn visit_mut_post<F>(&mut self, f: &mut F)
280 where
281 F: FnMut(&mut Self),
282 {
283 // This code uses `unsafe`. The core safety argument is that:
284 //
285 // - `children_mut()` produces disjoint children
286 // - no aliasing means each `Enter` is processed separately, and we `Leave` each node exactly once
287 //
288 // Put another way, our `stack` mirrors the function call stack, which allows multiple `&mut` refs at once,
289 // since only one stack frame can be active at a time.
290
291 use VisitMutAction::*;
292 let mut stack = vec![Enter(self as *mut T)];
293 while let Some(action) = stack.pop() {
294 match action {
295 Enter(ptr) => {
296 stack.push(Leave(ptr));
297 let elt = unsafe { &mut *ptr };
298 // Push children in reverse so they pop (and are visited) left-to-right.
299 stack.extend(elt.children_mut().rev().map(|child| Enter(child as *mut T)));
300 }
301 Leave(elt) => f(unsafe { &mut *elt }),
302 }
303 }
304 }
305
306 fn try_visit_post<'a, F, E>(&'a self, f: &mut F) -> Result<(), E>
307 where
308 F: FnMut(&'a Self) -> Result<(), E>,
309 {
310 use VisitAction::*;
311 let mut stack = vec![Enter(self)];
312 while let Some(action) = stack.pop() {
313 match action {
314 Enter(elt) => {
315 stack.push(Leave(elt));
316 // Push children in reverse so they pop (and are visited) left-to-right.
317 stack.extend(elt.children().rev().map(Enter));
318 }
319 Leave(elt) => f(elt)?,
320 }
321 }
322
323 Ok(())
324 }
325
326 #[allow(clippy::as_conversions)]
327 fn try_visit_mut_post<F, E>(&mut self, f: &mut F) -> Result<(), E>
328 where
329 F: FnMut(&mut Self) -> Result<(), E>,
330 {
331 // This code uses `unsafe`. The core safety argument is that:
332 //
333 // - `children_mut()` produces disjoint children
334 // - no aliasing means each `Enter` is processed separately, and we `Leave` each node exactly once
335 //
336 // Put another way, our `stack` mirrors the function call stack, which allows multiple `&mut` refs at once,
337 // since only one stack frame can be active at a time.
338
339 use VisitMutAction::*;
340 let mut stack = vec![Enter(self as *mut T)];
341 while let Some(action) = stack.pop() {
342 match action {
343 Enter(ptr) => {
344 stack.push(Leave(ptr));
345 let elt = unsafe { &mut *ptr };
346 // Push children in reverse so they pop (and are visited) left-to-right.
347 stack.extend(elt.children_mut().rev().map(|child| Enter(child as *mut T)));
348 }
349 Leave(ptr) => f(unsafe { &mut *ptr })?,
350 }
351 }
352
353 Ok(())
354 }
355
356 fn visit_pre<'a, F>(&'a self, f: &mut F)
357 where
358 F: FnMut(&'a Self),
359 {
360 let mut stack = vec![self];
361 while let Some(elt) = stack.pop() {
362 f(elt);
363 // Push children in reverse so they pop (and are visited) left-to-right.
364 stack.extend(elt.children().rev());
365 }
366 }
367
368 fn visit_pre_with_context<'a, Context, AccFun, Visitor>(
369 &'a self,
370 init: Context,
371 acc_fun: &mut AccFun,
372 visitor: &mut Visitor,
373 ) where
374 Context: Clone,
375 AccFun: FnMut(Context, &'a Self) -> Context,
376 Visitor: FnMut(&Context, &'a Self),
377 {
378 let mut stack = vec![(self, init)];
379 while let Some((elt, ctx)) = stack.pop() {
380 visitor(&ctx, elt);
381 let ctx = acc_fun(ctx, elt);
382 // Push children in reverse so they pop (and are visited) left-to-right.
383 stack.extend(elt.children().rev().map(|child| (child, ctx.clone())));
384 }
385 }
386
387 fn visit_mut_pre<F>(&mut self, f: &mut F)
388 where
389 F: FnMut(&mut Self),
390 {
391 let mut stack = vec![self];
392 while let Some(elt) = stack.pop() {
393 f(elt);
394 // Push children in reverse so they pop (and are visited) left-to-right.
395 stack.extend(elt.children_mut().rev())
396 }
397 }
398
399 fn try_visit_pre<'a, F, E>(&'a self, f: &mut F) -> Result<(), E>
400 where
401 F: FnMut(&'a Self) -> Result<(), E>,
402 {
403 let mut stack = vec![self];
404 while let Some(elt) = stack.pop() {
405 f(elt)?;
406 // Push children in reverse so they pop (and are visited) left-to-right.
407 stack.extend(elt.children().rev());
408 }
409
410 Ok(())
411 }
412
413 fn try_visit_mut_pre<F, E>(&mut self, f: &mut F) -> Result<(), E>
414 where
415 F: FnMut(&mut Self) -> Result<(), E>,
416 {
417 let mut stack = vec![self];
418 while let Some(elt) = stack.pop() {
419 f(elt)?;
420 // Push children in reverse so they pop (and are visited) left-to-right.
421 stack.extend(elt.children_mut().rev());
422 }
423
424 Ok(())
425 }
426 fn visit_pre_post<'a, F1, F2>(&'a self, pre: &mut F1, post: &mut F2)
427 where
428 F1: FnMut(&'a Self) -> Option<Vec<&'a Self>>,
429 F2: FnMut(&'a Self),
430 {
431 use VisitAction::*;
432 let mut stack = vec![Enter(self)];
433 while let Some(action) = stack.pop() {
434 match action {
435 Enter(elt) => {
436 stack.push(Leave(elt));
437 if let Some(children) = pre(elt) {
438 for child in children.into_iter().rev() {
439 stack.push(Enter(child));
440 }
441 } else {
442 for child in elt.children().rev() {
443 stack.push(Enter(child));
444 }
445 }
446 }
447 Leave(elt) => {
448 post(elt);
449 }
450 }
451 }
452 }
453
454 #[allow(clippy::as_conversions)]
455 fn visit_mut_pre_post<F1, F2>(&mut self, pre: &mut F1, post: &mut F2)
456 where
457 F1: FnMut(&mut Self) -> Option<Vec<&mut Self>>,
458 F2: FnMut(&mut Self),
459 {
460 // This code uses `unsafe`. The core safety argument is that:
461 //
462 // - `children_mut()` produces disjoint children
463 // - no aliasing means each `Enter` is processed separately, and we `Leave` each node exactly once
464 // - even if `pre` modifies the pointer, we retake it before computing children
465 //
466 // Put another way, our `stack` mirrors the function call stack, which allows multiple `&mut` refs at once,
467 // since only one stack frame can be active at a time.
468
469 use VisitMutAction::*;
470 let mut stack = vec![Enter(self as *mut T)];
471 while let Some(action) = stack.pop() {
472 match action {
473 Enter(ptr) => {
474 let elt = unsafe { &mut *ptr };
475 stack.push(Leave(ptr));
476
477 if let Some(children) = pre(elt) {
478 for child in children.into_iter().rev() {
479 stack.push(Enter(child));
480 }
481 } else {
482 let elt = unsafe { &mut *ptr };
483 for child in elt.children_mut().rev() {
484 stack.push(Enter(child));
485 }
486 }
487 }
488 Leave(ptr) => {
489 post(unsafe { &mut *ptr });
490 }
491 }
492 }
493 }
494}
495
496#[cfg(test)]
497mod tests {
498 use super::*;
499
500 // This test demonstrates how to build visitors for mutually recursive definitions.
501 // The key move here is the `direct_sub_*` methods, which are worklist-based traversals
502 // that find children of appropriate type.
503
504 #[derive(Debug, Eq, PartialEq)]
505 enum A {
506 Add(Box<A>, Box<A>),
507 Lit(u64),
508 FrB(Box<B>),
509 }
510
511 #[derive(Debug, Eq, PartialEq)]
512 enum B {
513 Mul(Box<B>, Box<B>),
514 Lit(u64),
515 FrA(Box<A>),
516 }
517
518 impl A {
519 fn direct_sub_b(&self) -> Vec<&B> {
520 let mut subs: Vec<&B> = vec![];
521
522 let mut worklist = vec![self];
523 while let Some(a) = worklist.pop() {
524 match a {
525 A::Add(lhs, rhs) => {
526 worklist.push(&*lhs);
527 worklist.push(&*rhs);
528 }
529 A::Lit(_) => (),
530 A::FrB(b) => subs.push(&*b),
531 }
532 }
533
534 subs
535 }
536
537 fn direct_sub_b_mut(&mut self) -> Vec<&mut B> {
538 let mut subs: Vec<&mut B> = vec![];
539
540 let mut worklist = vec![self];
541 while let Some(a) = worklist.pop() {
542 match a {
543 A::Add(lhs, rhs) => {
544 worklist.push(&mut **lhs);
545 worklist.push(&mut **rhs);
546 }
547 A::Lit(_) => (),
548 A::FrB(b) => subs.push(&mut **b),
549 }
550 }
551
552 subs
553 }
554 }
555
556 impl B {
557 fn direct_sub_a(&self) -> Vec<&A> {
558 let mut subs: Vec<&A> = vec![];
559
560 let mut worklist = vec![self];
561 while let Some(b) = worklist.pop() {
562 match b {
563 B::Mul(lhs, rhs) => {
564 worklist.push(&*lhs);
565 worklist.push(&*rhs);
566 }
567 B::Lit(_) => (),
568 B::FrA(a) => subs.push(&*a),
569 }
570 }
571
572 subs
573 }
574
575 fn direct_sub_a_mut(&mut self) -> Vec<&mut A> {
576 let mut subs: Vec<&mut A> = vec![];
577
578 let mut worklist = vec![self];
579 while let Some(b) = worklist.pop() {
580 match b {
581 B::Mul(lhs, rhs) => {
582 worklist.push(&mut **lhs);
583 worklist.push(&mut **rhs);
584 }
585 B::Lit(_) => (),
586 B::FrA(a) => subs.push(&mut **a),
587 }
588 }
589
590 subs
591 }
592 }
593
594 impl VisitChildren<A> for A {
595 fn visit_children<F>(&self, mut f: F)
596 where
597 F: FnMut(&A),
598 {
599 VisitChildren::visit_children(self, |expr: &B| {
600 Visit::visit_post(expr, &mut |expr| match expr {
601 B::FrA(expr) => f(expr.as_ref()),
602 _ => (),
603 });
604 });
605
606 match self {
607 A::Add(lhs, rhs) => {
608 f(lhs);
609 f(rhs);
610 }
611 A::Lit(_) => (),
612 A::FrB(_) => (),
613 }
614 }
615
616 fn visit_mut_children<F>(&mut self, mut f: F)
617 where
618 F: FnMut(&mut A),
619 {
620 VisitChildren::visit_mut_children(self, |expr: &mut B| {
621 Visit::visit_mut_post(expr, &mut |expr| match expr {
622 B::FrA(expr) => f(expr.as_mut()),
623 _ => (),
624 });
625 });
626
627 match self {
628 A::Add(lhs, rhs) => {
629 f(lhs);
630 f(rhs);
631 }
632 A::Lit(_) => (),
633 A::FrB(_) => (),
634 }
635 }
636
637 fn try_visit_children<F, E>(&self, mut f: F) -> Result<(), E>
638 where
639 F: FnMut(&A) -> Result<(), E>,
640 {
641 VisitChildren::try_visit_children(self, |expr: &B| {
642 Visit::try_visit_post(expr, &mut |expr| match expr {
643 B::FrA(expr) => f(expr.as_ref()),
644 _ => Ok(()),
645 })
646 })?;
647
648 match self {
649 A::Add(lhs, rhs) => {
650 f(lhs)?;
651 f(rhs)?;
652 }
653 A::Lit(_) => (),
654 A::FrB(_) => (),
655 }
656 Ok(())
657 }
658
659 fn try_visit_mut_children<F, E>(&mut self, mut f: F) -> Result<(), E>
660 where
661 F: FnMut(&mut A) -> Result<(), E>,
662 {
663 VisitChildren::try_visit_mut_children(self, |expr: &mut B| {
664 Visit::try_visit_mut_post(expr, &mut |expr| match expr {
665 B::FrA(expr) => f(expr.as_mut()),
666 _ => Ok(()),
667 })
668 })?;
669
670 match self {
671 A::Add(lhs, rhs) => {
672 f(lhs)?;
673 f(rhs)?;
674 }
675 A::Lit(_) => (),
676 A::FrB(_) => (),
677 }
678 Ok(())
679 }
680
681 fn children<'a>(&'a self) -> impl DoubleEndedIterator<Item = &'a A>
682 where
683 A: 'a,
684 {
685 let mut v: Vec<&A> = vec![];
686 match self {
687 A::Add(lhs, rhs) => {
688 v.push(&*lhs);
689 v.push(&*rhs)
690 }
691 A::Lit(_) => (),
692 A::FrB(b) => {
693 v.append(&mut b.direct_sub_a());
694 }
695 }
696
697 v.into_iter()
698 }
699
700 fn children_mut<'a>(&'a mut self) -> impl DoubleEndedIterator<Item = &'a mut A>
701 where
702 A: 'a,
703 {
704 let mut v: Vec<&mut A> = vec![];
705
706 match self {
707 A::Add(lhs, rhs) => {
708 v.push(&mut **lhs);
709 v.push(&mut **rhs)
710 }
711 A::Lit(_) => (),
712 A::FrB(b) => {
713 v.append(&mut b.direct_sub_a_mut());
714 }
715 }
716
717 v.into_iter()
718 }
719 }
720
721 impl VisitChildren<B> for A {
722 fn visit_children<F>(&self, mut f: F)
723 where
724 F: FnMut(&B),
725 {
726 match self {
727 A::Add(_, _) => (),
728 A::Lit(_) => (),
729 A::FrB(expr) => f(expr),
730 }
731 }
732
733 fn visit_mut_children<F>(&mut self, mut f: F)
734 where
735 F: FnMut(&mut B),
736 {
737 match self {
738 A::Add(_, _) => (),
739 A::Lit(_) => (),
740 A::FrB(expr) => f(expr),
741 }
742 }
743
744 fn try_visit_children<F, E>(&self, mut f: F) -> Result<(), E>
745 where
746 F: FnMut(&B) -> Result<(), E>,
747 {
748 match self {
749 A::Add(_, _) => Ok(()),
750 A::Lit(_) => Ok(()),
751 A::FrB(expr) => f(expr),
752 }
753 }
754
755 fn try_visit_mut_children<F, E>(&mut self, mut f: F) -> Result<(), E>
756 where
757 F: FnMut(&mut B) -> Result<(), E>,
758 {
759 match self {
760 A::Add(_, _) => Ok(()),
761 A::Lit(_) => Ok(()),
762 A::FrB(expr) => f(expr),
763 }
764 }
765
766 fn children<'a>(&'a self) -> impl DoubleEndedIterator<Item = &'a B>
767 where
768 B: 'a,
769 {
770 let mut child: Option<&B> = None;
771 match self {
772 A::Add(_, _) | A::Lit(_) => (),
773 A::FrB(b) => child = Some(&*b),
774 }
775 child.into_iter()
776 }
777
778 fn children_mut<'a>(&'a mut self) -> impl DoubleEndedIterator<Item = &'a mut B>
779 where
780 B: 'a,
781 {
782 let mut child: Option<&mut B> = None;
783 match self {
784 A::Add(_, _) | A::Lit(_) => (),
785 A::FrB(b) => child = Some(&mut **b),
786 }
787 child.into_iter()
788 }
789 }
790
791 impl VisitChildren<B> for B {
792 fn visit_children<F>(&self, mut f: F)
793 where
794 F: FnMut(&B),
795 {
796 // VisitChildren::visit_children(self, |expr: &A| {
797 // #[allow(deprecated)]
798 // Visit::visit_post(expr, &mut |expr| match expr {
799 // A::FrB(expr) => f(expr.as_ref()),
800 // _ => (),
801 // });
802 // });
803
804 match self {
805 B::Mul(lhs, rhs) => {
806 f(lhs);
807 f(rhs);
808 }
809 B::Lit(_) => (),
810 B::FrA(_) => (),
811 }
812 }
813
814 fn visit_mut_children<F>(&mut self, mut f: F)
815 where
816 F: FnMut(&mut B),
817 {
818 // VisitChildren::visit_mut_children(self, |expr: &mut A| {
819 // #[allow(deprecated)]
820 // Visit::visit_mut_post_nolimit(expr, &mut |expr| match expr {
821 // A::FrB(expr) => f(expr.as_mut()),
822 // _ => (),
823 // });
824 // });
825
826 match self {
827 B::Mul(lhs, rhs) => {
828 f(lhs);
829 f(rhs);
830 }
831 B::Lit(_) => (),
832 B::FrA(_) => (),
833 }
834 }
835
836 fn try_visit_children<F, E>(&self, mut f: F) -> Result<(), E>
837 where
838 F: FnMut(&B) -> Result<(), E>,
839 {
840 // VisitChildren::try_visit_children(self, |expr: &A| {
841 // Visit::try_visit_post(expr, &mut |expr| match expr {
842 // A::FrB(expr) => f(expr.as_ref()),
843 // _ => Ok(()),
844 // })
845 // })?;
846
847 match self {
848 B::Mul(lhs, rhs) => {
849 f(lhs)?;
850 f(rhs)?;
851 }
852 B::Lit(_) => (),
853 B::FrA(_) => (),
854 }
855 Ok(())
856 }
857
858 fn try_visit_mut_children<F, E>(&mut self, mut f: F) -> Result<(), E>
859 where
860 F: FnMut(&mut B) -> Result<(), E>,
861 {
862 // VisitChildren::try_visit_mut_children(self, |expr: &mut A| {
863 // Visit::try_visit_mut_post(expr, &mut |expr| match expr {
864 // A::FrB(expr) => f(expr.as_mut()),
865 // _ => Ok(()),
866 // })
867 // })?;
868
869 match self {
870 B::Mul(lhs, rhs) => {
871 f(lhs)?;
872 f(rhs)?;
873 }
874 B::Lit(_) => (),
875 B::FrA(_) => (),
876 }
877 Ok(())
878 }
879
880 fn children<'a>(&'a self) -> impl DoubleEndedIterator<Item = &'a B>
881 where
882 B: 'a,
883 {
884 let mut v: Vec<&B> = vec![];
885 match self {
886 B::Mul(lhs, rhs) => {
887 v.push(&*lhs);
888 v.push(&*rhs);
889 }
890 B::Lit(_) => (),
891 B::FrA(a) => v.append(&mut a.direct_sub_b()),
892 }
893 v.into_iter()
894 }
895
896 fn children_mut<'a>(&'a mut self) -> impl DoubleEndedIterator<Item = &'a mut B>
897 where
898 B: 'a,
899 {
900 let mut v: Vec<&mut B> = vec![];
901 match self {
902 B::Mul(lhs, rhs) => {
903 v.push(&mut **lhs);
904 v.push(&mut **rhs);
905 }
906 B::Lit(_) => (),
907 B::FrA(a) => v.append(&mut a.direct_sub_b_mut()),
908 }
909 v.into_iter()
910 }
911 }
912
913 impl VisitChildren<A> for B {
914 fn visit_children<F>(&self, mut f: F)
915 where
916 F: FnMut(&A),
917 {
918 match self {
919 B::Mul(_, _) => (),
920 B::Lit(_) => (),
921 B::FrA(expr) => f(expr),
922 }
923 }
924
925 fn visit_mut_children<F>(&mut self, mut f: F)
926 where
927 F: FnMut(&mut A),
928 {
929 match self {
930 B::Mul(_, _) => (),
931 B::Lit(_) => (),
932 B::FrA(expr) => f(expr),
933 }
934 }
935
936 fn try_visit_children<F, E>(&self, mut f: F) -> Result<(), E>
937 where
938 F: FnMut(&A) -> Result<(), E>,
939 {
940 match self {
941 B::Mul(_, _) => Ok(()),
942 B::Lit(_) => Ok(()),
943 B::FrA(expr) => f(expr),
944 }
945 }
946
947 fn try_visit_mut_children<F, E>(&mut self, mut f: F) -> Result<(), E>
948 where
949 F: FnMut(&mut A) -> Result<(), E>,
950 {
951 match self {
952 B::Mul(_, _) => Ok(()),
953 B::Lit(_) => Ok(()),
954 B::FrA(expr) => f(expr),
955 }
956 }
957
958 fn children<'a>(&'a self) -> impl DoubleEndedIterator<Item = &'a A>
959 where
960 A: 'a,
961 {
962 let mut child: Option<&A> = None;
963 match self {
964 B::Mul(_, _) | B::Lit(_) => (),
965 B::FrA(a) => child = Some(&*a),
966 }
967 child.into_iter()
968 }
969
970 fn children_mut<'a>(&'a mut self) -> impl DoubleEndedIterator<Item = &'a mut A>
971 where
972 A: 'a,
973 {
974 let mut child: Option<&mut A> = None;
975 match self {
976 B::Mul(_, _) | B::Lit(_) => (),
977 B::FrA(a) => child = Some(&mut **a),
978 }
979 child.into_iter()
980 }
981 }
982
983 /// x + (y + z)
984 fn test_term_a(x: A, y: A, z: A) -> A {
985 let x = Box::new(x);
986 let y = Box::new(y);
987 let z = Box::new(z);
988 A::Add(x, Box::new(A::Add(y, z)))
989 }
990
991 /// u + (v + w)
992 fn test_term_b(u: B, v: B, w: B) -> B {
993 let u = Box::new(u);
994 let v = Box::new(v);
995 let w = Box::new(w);
996 B::Mul(u, Box::new(B::Mul(v, w)))
997 }
998
999 fn a_to_b(x: A) -> B {
1000 B::FrA(Box::new(x))
1001 }
1002
1003 fn b_to_a(x: B) -> A {
1004 A::FrB(Box::new(x))
1005 }
1006
1007 fn test_term_rec_b(b: u64) -> B {
1008 test_term_b(
1009 a_to_b(test_term_a(
1010 b_to_a(test_term_b(B::Lit(b + 11), B::Lit(b + 12), B::Lit(b + 13))),
1011 b_to_a(test_term_b(B::Lit(b + 14), B::Lit(b + 15), B::Lit(b + 16))),
1012 b_to_a(test_term_b(B::Lit(b + 17), B::Lit(b + 18), B::Lit(b + 19))),
1013 )),
1014 a_to_b(test_term_a(
1015 b_to_a(test_term_b(B::Lit(b + 21), B::Lit(b + 22), B::Lit(b + 23))),
1016 b_to_a(test_term_b(B::Lit(b + 24), B::Lit(b + 25), B::Lit(b + 26))),
1017 b_to_a(test_term_b(B::Lit(b + 27), B::Lit(b + 28), B::Lit(b + 29))),
1018 )),
1019 a_to_b(test_term_a(
1020 b_to_a(test_term_b(B::Lit(b + 31), B::Lit(b + 32), B::Lit(b + 33))),
1021 b_to_a(test_term_b(B::Lit(b + 34), B::Lit(b + 35), B::Lit(b + 36))),
1022 b_to_a(test_term_b(B::Lit(b + 37), B::Lit(b + 38), B::Lit(b + 39))),
1023 )),
1024 )
1025 }
1026
1027 fn test_term_rec_a(b: u64) -> A {
1028 test_term_a(
1029 b_to_a(test_term_b(
1030 a_to_b(test_term_a(A::Lit(b + 11), A::Lit(b + 12), A::Lit(b + 13))),
1031 a_to_b(test_term_a(A::Lit(b + 14), A::Lit(b + 15), A::Lit(b + 16))),
1032 a_to_b(test_term_a(A::Lit(b + 17), A::Lit(b + 18), A::Lit(b + 19))),
1033 )),
1034 b_to_a(test_term_b(
1035 a_to_b(test_term_a(A::Lit(b + 21), A::Lit(b + 22), A::Lit(b + 23))),
1036 a_to_b(test_term_a(A::Lit(b + 24), A::Lit(b + 25), A::Lit(b + 26))),
1037 a_to_b(test_term_a(A::Lit(b + 27), A::Lit(b + 28), A::Lit(b + 29))),
1038 )),
1039 b_to_a(test_term_b(
1040 a_to_b(test_term_a(A::Lit(b + 31), A::Lit(b + 32), A::Lit(b + 33))),
1041 a_to_b(test_term_a(A::Lit(b + 34), A::Lit(b + 35), A::Lit(b + 36))),
1042 a_to_b(test_term_a(A::Lit(b + 37), A::Lit(b + 38), A::Lit(b + 39))),
1043 )),
1044 )
1045 }
1046
1047 #[mz_ore::test]
1048 #[cfg_attr(miri, ignore)] // error: unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux`
1049 fn test_recursive_types_a() {
1050 let mut act = test_term_rec_a(0);
1051 let exp = test_term_rec_a(20);
1052
1053 act.visit_mut_pre(&mut |expr| match expr {
1054 A::Lit(x) => *x = *x + 20,
1055 _ => (),
1056 });
1057
1058 assert_eq!(act, exp);
1059 }
1060
1061 #[mz_ore::test]
1062 fn test_recursive_types_b() {
1063 let mut act = test_term_rec_b(0);
1064 let exp = test_term_rec_b(30);
1065
1066 act.visit_mut_pre(&mut |expr| match expr {
1067 B::Lit(x) => *x = *x + 30,
1068 _ => (),
1069 });
1070
1071 assert_eq!(act, exp);
1072 }
1073}