Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/common/src/tree_node.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
//! [`TreeNode`] for visiting and rewriting expression and plan trees
19
20
use std::sync::Arc;
21
22
use crate::Result;
23
24
/// These macros are used to determine continuation during transforming traversals.
25
macro_rules! handle_transform_recursion {
26
    ($F_DOWN:expr, $F_CHILD:expr, $F_UP:expr) => {{
27
        $F_DOWN?
28
0
            .transform_children(|n| n.map_children($F_CHILD))?
29
            .transform_parent($F_UP)
30
    }};
31
}
32
33
/// API for inspecting and rewriting tree data structures.
34
///
35
/// The `TreeNode` API is used to express algorithms separately from traversing
36
/// the structure of `TreeNode`s, avoiding substantial code duplication.
37
///
38
/// This trait is implemented for plans ([`ExecutionPlan`], [`LogicalPlan`]) and
39
/// expression trees ([`PhysicalExpr`], [`Expr`]) as well as Plan+Payload
40
/// combinations [`PlanContext`] and [`ExprContext`].
41
///
42
/// # Overview
43
/// There are three categories of TreeNode APIs:
44
///
45
/// 1. "Inspecting" APIs to traverse a tree of `&TreeNodes`:
46
///    [`apply`], [`visit`], [`exists`].
47
///
48
/// 2. "Transforming" APIs that traverse and consume a tree of `TreeNode`s
49
///    producing possibly changed `TreeNode`s: [`transform`], [`transform_up`],
50
///    [`transform_down`], [`transform_down_up`], and [`rewrite`].
51
///
52
/// 3. Internal APIs used to implement the `TreeNode` API: [`apply_children`],
53
///    and [`map_children`].
54
///
55
/// | Traversal Order | Inspecting | Transforming |
56
/// | --- | --- | --- |
57
/// | top-down | [`apply`], [`exists`] | [`transform_down`]|
58
/// | bottom-up | | [`transform`] , [`transform_up`]|
59
/// | combined with separate `f_down` and `f_up` closures | | [`transform_down_up`] |
60
/// | combined with `f_down()` and `f_up()` in an object | [`visit`]  | [`rewrite`] |
61
///
62
/// **Note**:while there is currently no in-place mutation API that uses `&mut
63
/// TreeNode`, the transforming APIs are efficient and optimized to avoid
64
/// cloning.
65
///
66
/// [`apply`]: Self::apply
67
/// [`visit`]: Self::visit
68
/// [`exists`]: Self::exists
69
/// [`transform`]: Self::transform
70
/// [`transform_up`]: Self::transform_up
71
/// [`transform_down`]: Self::transform_down
72
/// [`transform_down_up`]: Self::transform_down_up
73
/// [`rewrite`]: Self::rewrite
74
/// [`apply_children`]: Self::apply_children
75
/// [`map_children`]: Self::map_children
76
///
77
/// # Terminology
78
/// The following terms are used in this trait
79
///
80
/// * `f_down`: Invoked before any children of the current node are visited.
81
/// * `f_up`: Invoked after all children of the current node are visited.
82
/// * `f`: closure that is applied to the current node.
83
/// * `map_*`: applies a transformation to rewrite owned nodes
84
/// * `apply_*`:  invokes a function on borrowed nodes
85
/// * `transform_`: applies a transformation to rewrite owned nodes
86
///
87
/// <!-- Since these are in the datafusion-common crate, can't use intra doc links) -->
88
/// [`ExecutionPlan`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.ExecutionPlan.html
89
/// [`PhysicalExpr`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.PhysicalExpr.html
90
/// [`LogicalPlan`]: https://docs.rs/datafusion-expr/latest/datafusion_expr/logical_plan/enum.LogicalPlan.html
91
/// [`Expr`]: https://docs.rs/datafusion-expr/latest/datafusion_expr/expr/enum.Expr.html
92
/// [`PlanContext`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/tree_node/struct.PlanContext.html
93
/// [`ExprContext`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/tree_node/struct.ExprContext.html
94
pub trait TreeNode: Sized {
95
    /// Visit the tree node with a [`TreeNodeVisitor`], performing a
96
    /// depth-first walk of the node and its children.
97
    ///
98
    /// [`TreeNodeVisitor::f_down()`] is called in top-down order (before
99
    /// children are visited), [`TreeNodeVisitor::f_up()`] is called in
100
    /// bottom-up order (after children are visited).
101
    ///
102
    /// # Return Value
103
    /// Specifies how the tree walk ended. See [`TreeNodeRecursion`] for details.
104
    ///
105
    /// # See Also:
106
    /// * [`Self::apply`] for inspecting nodes with a closure
107
    /// * [`Self::rewrite`] to rewrite owned `TreeNode`s
108
    ///
109
    /// # Example
110
    /// Consider the following tree structure:
111
    /// ```text
112
    /// ParentNode
113
    ///    left: ChildNode1
114
    ///    right: ChildNode2
115
    /// ```
116
    ///
117
    /// Here, the nodes would be visited using the following order:
118
    /// ```text
119
    /// TreeNodeVisitor::f_down(ParentNode)
120
    /// TreeNodeVisitor::f_down(ChildNode1)
121
    /// TreeNodeVisitor::f_up(ChildNode1)
122
    /// TreeNodeVisitor::f_down(ChildNode2)
123
    /// TreeNodeVisitor::f_up(ChildNode2)
124
    /// TreeNodeVisitor::f_up(ParentNode)
125
    /// ```
126
0
    fn visit<'n, V: TreeNodeVisitor<'n, Node = Self>>(
127
0
        &'n self,
128
0
        visitor: &mut V,
129
0
    ) -> Result<TreeNodeRecursion> {
130
0
        visitor
131
0
            .f_down(self)?
132
0
            .visit_children(|| self.apply_children(|c| c.visit(visitor)))?
133
0
            .visit_parent(|| visitor.f_up(self))
134
0
    }
135
136
    /// Rewrite the tree node with a [`TreeNodeRewriter`], performing a
137
    /// depth-first walk of the node and its children.
138
    ///
139
    /// [`TreeNodeRewriter::f_down()`] is called in top-down order (before
140
    /// children are visited), [`TreeNodeRewriter::f_up()`] is called in
141
    /// bottom-up order (after children are visited).
142
    ///
143
    /// Note: If using the default [`TreeNodeRewriter::f_up`] or
144
    /// [`TreeNodeRewriter::f_down`] that do nothing, consider using
145
    /// [`Self::transform_down`] instead.
146
    ///
147
    /// # Return Value
148
    /// The returns value specifies how the tree walk should proceed. See
149
    /// [`TreeNodeRecursion`] for details. If an [`Err`] is returned, the
150
    /// recursion stops immediately.
151
    ///
152
    /// # See Also
153
    /// * [`Self::visit`] for inspecting (without modification) `TreeNode`s
154
    /// * [Self::transform_down_up] for a top-down (pre-order) traversal.
155
    /// * [Self::transform_down] for a top-down (pre-order) traversal.
156
    /// * [`Self::transform_up`] for a bottom-up (post-order) traversal.
157
    ///
158
    /// # Example
159
    /// Consider the following tree structure:
160
    /// ```text
161
    /// ParentNode
162
    ///    left: ChildNode1
163
    ///    right: ChildNode2
164
    /// ```
165
    ///
166
    /// Here, the nodes would be visited using the following order:
167
    /// ```text
168
    /// TreeNodeRewriter::f_down(ParentNode)
169
    /// TreeNodeRewriter::f_down(ChildNode1)
170
    /// TreeNodeRewriter::f_up(ChildNode1)
171
    /// TreeNodeRewriter::f_down(ChildNode2)
172
    /// TreeNodeRewriter::f_up(ChildNode2)
173
    /// TreeNodeRewriter::f_up(ParentNode)
174
    /// ```
175
0
    fn rewrite<R: TreeNodeRewriter<Node = Self>>(
176
0
        self,
177
0
        rewriter: &mut R,
178
0
    ) -> Result<Transformed<Self>> {
179
0
        handle_transform_recursion!(rewriter.f_down(self), |c| c.rewrite(rewriter), |n| {
180
0
            rewriter.f_up(n)
181
0
        })
182
0
    }
183
184
    /// Applies `f` to the node then each of its children, recursively (a
185
    /// top-down, pre-order traversal).
186
    ///
187
    /// The return [`TreeNodeRecursion`] controls the recursion and can cause
188
    /// an early return.
189
    ///
190
    /// # See Also
191
    /// * [`Self::transform_down`] for the equivalent transformation API.
192
    /// * [`Self::visit`] for both top-down and bottom up traversal.
193
4.50k
    fn apply<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
194
4.50k
        &'n self,
195
4.50k
        mut f: F,
196
4.50k
    ) -> Result<TreeNodeRecursion> {
197
4.90k
        fn apply_impl<'n, N: TreeNode, F: FnMut(&'n N) -> Result<TreeNodeRecursion>>(
198
4.90k
            node: &'n N,
199
4.90k
            f: &mut F,
200
4.90k
        ) -> Result<TreeNodeRecursion> {
201
4.90k
            f(node)
?0
.visit_children(|| node.apply_children(|c|
apply_impl(c, f)402
))
202
4.90k
        }
203
204
4.50k
        apply_impl(self, &mut f)
205
4.50k
    }
206
207
    /// Recursively rewrite the node's children and then the node using `f`
208
    /// (a bottom-up post-order traversal).
209
    ///
210
    /// A synonym of [`Self::transform_up`].
211
6.62k
    fn transform<F: FnMut(Self) -> Result<Transformed<Self>>>(
212
6.62k
        self,
213
6.62k
        f: F,
214
6.62k
    ) -> Result<Transformed<Self>> {
215
6.62k
        self.transform_up(f)
216
6.62k
    }
217
218
    /// Recursively rewrite the tree using `f` in a top-down (pre-order)
219
    /// fashion.
220
    ///
221
    /// `f` is applied to the node first, and then its children.
222
    ///
223
    /// # See Also
224
    /// * [`Self::transform_up`] for a bottom-up (post-order) traversal.
225
    /// * [Self::transform_down_up] for a combined traversal with closures
226
    /// * [`Self::rewrite`] for a combined traversal with a visitor
227
363
    fn transform_down<F: FnMut(Self) -> Result<Transformed<Self>>>(
228
363
        self,
229
363
        mut f: F,
230
363
    ) -> Result<Transformed<Self>> {
231
363
        fn transform_down_impl<N: TreeNode, F: FnMut(N) -> Result<Transformed<N>>>(
232
363
            node: N,
233
363
            f: &mut F,
234
363
        ) -> Result<Transformed<N>> {
235
363
            f(node)
?0
.transform_children(|n| n.map_children(|c|
transform_down_impl(c, f)0
))
236
363
        }
237
238
363
        transform_down_impl(self, &mut f)
239
363
    }
240
241
    /// Same as [`Self::transform_down`] but with a mutable closure.
242
    #[deprecated(since = "38.0.0", note = "Use `transform_down` instead")]
243
0
    fn transform_down_mut<F: FnMut(Self) -> Result<Transformed<Self>>>(
244
0
        self,
245
0
        f: &mut F,
246
0
    ) -> Result<Transformed<Self>> {
247
0
        self.transform_down(f)
248
0
    }
249
250
    /// Recursively rewrite the node using `f` in a bottom-up (post-order)
251
    /// fashion.
252
    ///
253
    /// `f` is applied to the node's  children first, and then to the node itself.
254
    ///
255
    /// # See Also
256
    /// * [`Self::transform_down`] top-down (pre-order) traversal.
257
    /// * [Self::transform_down_up] for a combined traversal with closures
258
    /// * [`Self::rewrite`] for a combined traversal with a visitor
259
10.4k
    fn transform_up<F: FnMut(Self) -> Result<Transformed<Self>>>(
260
10.4k
        self,
261
10.4k
        mut f: F,
262
10.4k
    ) -> Result<Transformed<Self>> {
263
26.5k
        fn transform_up_impl<N: TreeNode, F: FnMut(N) -> Result<Transformed<N>>>(
264
26.5k
            node: N,
265
26.5k
            f: &mut F,
266
26.5k
        ) -> Result<Transformed<N>> {
267
26.5k
            node.map_children(|c| 
transform_up_impl(c, f)16.0k
)
?0
268
26.5k
                .transform_parent(f)
269
26.5k
        }
270
271
10.4k
        transform_up_impl(self, &mut f)
272
10.4k
    }
273
274
    /// Same as [`Self::transform_up`] but with a mutable closure.
275
    #[deprecated(since = "38.0.0", note = "Use `transform_up` instead")]
276
0
    fn transform_up_mut<F: FnMut(Self) -> Result<Transformed<Self>>>(
277
0
        self,
278
0
        f: &mut F,
279
0
    ) -> Result<Transformed<Self>> {
280
0
        self.transform_up(f)
281
0
    }
282
283
    /// Transforms the node using `f_down` while traversing the tree top-down
284
    /// (pre-order), and using `f_up` while traversing the tree bottom-up
285
    /// (post-order).
286
    ///
287
    /// The method behaves the same as calling [`Self::transform_down`] followed
288
    /// by [`Self::transform_up`] on the same node. Use this method if you want
289
    /// to start the `f_up` process right where `f_down` jumps. This can make
290
    /// the whole process faster by reducing the number of `f_up` steps.
291
    ///
292
    /// # See Also
293
    /// * [`Self::transform_up`] for a bottom-up (post-order) traversal.
294
    /// * [Self::transform_down] for a top-down (pre-order) traversal.
295
    /// * [`Self::rewrite`] for a combined traversal with a visitor
296
    ///
297
    /// # Example
298
    /// Consider the following tree structure:
299
    /// ```text
300
    /// ParentNode
301
    ///    left: ChildNode1
302
    ///    right: ChildNode2
303
    /// ```
304
    ///
305
    /// The nodes are visited using the following order:
306
    /// ```text
307
    /// f_down(ParentNode)
308
    /// f_down(ChildNode1)
309
    /// f_up(ChildNode1)
310
    /// f_down(ChildNode2)
311
    /// f_up(ChildNode2)
312
    /// f_up(ParentNode)
313
    /// ```
314
    ///
315
    /// See [`TreeNodeRecursion`] for more details on controlling the traversal.
316
    ///
317
    /// If `f_down` or `f_up` returns [`Err`], the recursion stops immediately.
318
    ///
319
    /// Example:
320
    /// ```text
321
    ///                                               |   +---+
322
    ///                                               |   | J |
323
    ///                                               |   +---+
324
    ///                                               |     |
325
    ///                                               |   +---+
326
    ///                  TreeNodeRecursion::Continue  |   | I |
327
    ///                                               |   +---+
328
    ///                                               |     |
329
    ///                                               |   +---+
330
    ///                                              \|/  | F |
331
    ///                                               '   +---+
332
    ///                                                  /     \ ___________________
333
    ///                  When `f_down` is           +---+                           \ ---+
334
    ///                  applied on node "E",       | E |                            | G |
335
    ///                  it returns with "Jump".    +---+                            +---+
336
    ///                                               |                                |
337
    ///                                             +---+                            +---+
338
    ///                                             | C |                            | H |
339
    ///                                             +---+                            +---+
340
    ///                                             /   \
341
    ///                                        +---+     +---+
342
    ///                                        | B |     | D |
343
    ///                                        +---+     +---+
344
    ///                                                    |
345
    ///                                                  +---+
346
    ///                                                  | A |
347
    ///                                                  +---+
348
    ///
349
    /// Instead of starting from leaf nodes, `f_up` starts from the node "E".
350
    ///                                                   +---+
351
    ///                                               |   | J |
352
    ///                                               |   +---+
353
    ///                                               |     |
354
    ///                                               |   +---+
355
    ///                                               |   | I |
356
    ///                                               |   +---+
357
    ///                                               |     |
358
    ///                                              /    +---+
359
    ///                                            /      | F |
360
    ///                                          /        +---+
361
    ///                                        /         /     \ ______________________
362
    ///                                       |     +---+   .                          \ ---+
363
    ///                                       |     | E |  /|\  After `f_down` jumps    | G |
364
    ///                                       |     +---+   |   on node E, `f_up`       +---+
365
    ///                                        \------| ---/   if applied on node E.      |
366
    ///                                             +---+                               +---+
367
    ///                                             | C |                               | H |
368
    ///                                             +---+                               +---+
369
    ///                                             /   \
370
    ///                                        +---+     +---+
371
    ///                                        | B |     | D |
372
    ///                                        +---+     +---+
373
    ///                                                    |
374
    ///                                                  +---+
375
    ///                                                  | A |
376
    ///                                                  +---+
377
    /// ```
378
0
    fn transform_down_up<
379
0
        FD: FnMut(Self) -> Result<Transformed<Self>>,
380
0
        FU: FnMut(Self) -> Result<Transformed<Self>>,
381
0
    >(
382
0
        self,
383
0
        mut f_down: FD,
384
0
        mut f_up: FU,
385
0
    ) -> Result<Transformed<Self>> {
386
0
        fn transform_down_up_impl<
387
0
            N: TreeNode,
388
0
            FD: FnMut(N) -> Result<Transformed<N>>,
389
0
            FU: FnMut(N) -> Result<Transformed<N>>,
390
0
        >(
391
0
            node: N,
392
0
            f_down: &mut FD,
393
0
            f_up: &mut FU,
394
0
        ) -> Result<Transformed<N>> {
395
0
            handle_transform_recursion!(
396
0
                f_down(node),
397
0
                |c| transform_down_up_impl(c, f_down, f_up),
398
0
                f_up
399
            )
400
0
        }
401
402
0
        transform_down_up_impl(self, &mut f_down, &mut f_up)
403
0
    }
404
405
    /// Returns true if `f` returns true for any node in the tree.
406
    ///
407
    /// Stops recursion as soon as a matching node is found
408
0
    fn exists<F: FnMut(&Self) -> Result<bool>>(&self, mut f: F) -> Result<bool> {
409
0
        let mut found = false;
410
0
        self.apply(|n| {
411
0
            Ok(if f(n)? {
412
0
                found = true;
413
0
                TreeNodeRecursion::Stop
414
            } else {
415
0
                TreeNodeRecursion::Continue
416
            })
417
0
        })
418
0
        .map(|_| found)
419
0
    }
420
421
    /// Low-level API used to implement other APIs.
422
    ///
423
    /// If you want to implement the [`TreeNode`] trait for your own type, you
424
    /// should implement this method and [`Self::map_children`].
425
    ///
426
    /// Users should use one of the higher level APIs described on [`Self`].
427
    ///
428
    /// Description: Apply `f` to inspect node's children (but not the node
429
    /// itself).
430
    fn apply_children<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
431
        &'n self,
432
        f: F,
433
    ) -> Result<TreeNodeRecursion>;
434
435
    /// Low-level API used to implement other APIs.
436
    ///
437
    /// If you want to implement the [`TreeNode`] trait for your own type, you
438
    /// should implement this method and [`Self::apply_children`].
439
    ///
440
    /// Users should use one of the higher level APIs described on [`Self`].
441
    ///
442
    /// Description: Apply `f` to rewrite the node's children (but not the node itself).
443
    fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>(
444
        self,
445
        f: F,
446
    ) -> Result<Transformed<Self>>;
447
}
448
449
/// A [Visitor](https://en.wikipedia.org/wiki/Visitor_pattern) for recursively
450
/// inspecting [`TreeNode`]s via [`TreeNode::visit`].
451
///
452
/// See [`TreeNode`] for more details on available APIs
453
///
454
/// When passed to [`TreeNode::visit`], [`TreeNodeVisitor::f_down`] and
455
/// [`TreeNodeVisitor::f_up`] are invoked recursively on the tree.
456
/// See [`TreeNodeRecursion`] for more details on controlling the traversal.
457
///
458
/// # Return Value
459
/// The returns value of `f_up` and `f_down` specifies how the tree walk should
460
/// proceed. See [`TreeNodeRecursion`] for details. If an [`Err`] is returned,
461
/// the recursion stops immediately.
462
///
463
/// Note: If using the default implementations of [`TreeNodeVisitor::f_up`] or
464
/// [`TreeNodeVisitor::f_down`] that do nothing, consider using
465
/// [`TreeNode::apply`] instead.
466
///
467
/// # See Also:
468
/// * [`TreeNode::rewrite`] to rewrite owned `TreeNode`s
469
pub trait TreeNodeVisitor<'n>: Sized {
470
    /// The node type which is visitable.
471
    type Node: TreeNode;
472
473
    /// Invoked while traversing down the tree, before any children are visited.
474
    /// Default implementation continues the recursion.
475
0
    fn f_down(&mut self, _node: &'n Self::Node) -> Result<TreeNodeRecursion> {
476
0
        Ok(TreeNodeRecursion::Continue)
477
0
    }
478
479
    /// Invoked while traversing up the tree after children are visited. Default
480
    /// implementation continues the recursion.
481
0
    fn f_up(&mut self, _node: &'n Self::Node) -> Result<TreeNodeRecursion> {
482
0
        Ok(TreeNodeRecursion::Continue)
483
0
    }
484
}
485
486
/// A [Visitor](https://en.wikipedia.org/wiki/Visitor_pattern) for recursively
487
/// rewriting [`TreeNode`]s via [`TreeNode::rewrite`].
488
///
489
/// For example you can implement this trait on a struct to rewrite `Expr` or
490
/// `LogicalPlan` that needs to track state during the rewrite.
491
///
492
/// See [`TreeNode`] for more details on available APIs
493
///
494
/// When passed to [`TreeNode::rewrite`], [`TreeNodeRewriter::f_down`] and
495
/// [`TreeNodeRewriter::f_up`] are invoked recursively on the tree.
496
/// See [`TreeNodeRecursion`] for more details on controlling the traversal.
497
///
498
/// # Return Value
499
/// The returns value of `f_up` and `f_down` specifies how the tree walk should
500
/// proceed. See [`TreeNodeRecursion`] for details. If an [`Err`] is returned,
501
/// the recursion stops immediately.
502
///
503
/// Note: If using the default implementations of [`TreeNodeRewriter::f_up`] or
504
/// [`TreeNodeRewriter::f_down`] that do nothing, consider using
505
/// [`TreeNode::transform_up`] or [`TreeNode::transform_down`] instead.
506
///
507
/// # See Also:
508
/// * [`TreeNode::visit`] to inspect borrowed `TreeNode`s
509
pub trait TreeNodeRewriter: Sized {
510
    /// The node type which is rewritable.
511
    type Node: TreeNode;
512
513
    /// Invoked while traversing down the tree before any children are rewritten.
514
    /// Default implementation returns the node as is and continues recursion.
515
0
    fn f_down(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
516
0
        Ok(Transformed::no(node))
517
0
    }
518
519
    /// Invoked while traversing up the tree after all children have been rewritten.
520
    /// Default implementation returns the node as is and continues recursion.
521
0
    fn f_up(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
522
0
        Ok(Transformed::no(node))
523
0
    }
524
}
525
526
/// Controls how [`TreeNode`] recursions should proceed.
527
#[derive(Debug, PartialEq, Clone, Copy)]
528
pub enum TreeNodeRecursion {
529
    /// Continue recursion with the next node.
530
    Continue,
531
    /// In top-down traversals, skip recursing into children but continue with
532
    /// the next node, which actually means pruning of the subtree.
533
    ///
534
    /// In bottom-up traversals, bypass calling bottom-up closures till the next
535
    /// leaf node.
536
    ///
537
    /// In combined traversals, if it is the `f_down` (pre-order) phase, execution
538
    /// "jumps" to the next `f_up` (post-order) phase by shortcutting its children.
539
    /// If it is the `f_up` (post-order) phase, execution "jumps" to the next `f_down`
540
    /// (pre-order) phase by shortcutting its parent nodes until the first parent node
541
    /// having unvisited children path.
542
    Jump,
543
    /// Stop recursion.
544
    Stop,
545
}
546
547
impl TreeNodeRecursion {
548
    /// Continues visiting nodes with `f` depending on the current [`TreeNodeRecursion`]
549
    /// value and the fact that `f` is visiting the current node's children.
550
4.90k
    pub fn visit_children<F: FnOnce() -> Result<TreeNodeRecursion>>(
551
4.90k
        self,
552
4.90k
        f: F,
553
4.90k
    ) -> Result<TreeNodeRecursion> {
554
4.90k
        match self {
555
4.90k
            TreeNodeRecursion::Continue => f(),
556
0
            TreeNodeRecursion::Jump => Ok(TreeNodeRecursion::Continue),
557
0
            TreeNodeRecursion::Stop => Ok(self),
558
        }
559
4.90k
    }
560
561
    /// Continues visiting nodes with `f` depending on the current [`TreeNodeRecursion`]
562
    /// value and the fact that `f` is visiting the current node's sibling.
563
0
    pub fn visit_sibling<F: FnOnce() -> Result<TreeNodeRecursion>>(
564
0
        self,
565
0
        f: F,
566
0
    ) -> Result<TreeNodeRecursion> {
567
0
        match self {
568
0
            TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => f(),
569
0
            TreeNodeRecursion::Stop => Ok(self),
570
        }
571
0
    }
572
573
    /// Continues visiting nodes with `f` depending on the current [`TreeNodeRecursion`]
574
    /// value and the fact that `f` is visiting the current node's parent.
575
0
    pub fn visit_parent<F: FnOnce() -> Result<TreeNodeRecursion>>(
576
0
        self,
577
0
        f: F,
578
0
    ) -> Result<TreeNodeRecursion> {
579
0
        match self {
580
0
            TreeNodeRecursion::Continue => f(),
581
0
            TreeNodeRecursion::Jump | TreeNodeRecursion::Stop => Ok(self),
582
        }
583
0
    }
584
}
585
586
/// Result of tree walk / transformation APIs
587
///
588
/// `Transformed` is a wrapper around the tree node data (e.g. `Expr` or
589
/// `LogicalPlan`). It is used to indicate whether the node was transformed
590
/// and how the recursion should proceed.
591
///
592
/// [`TreeNode`] API users control the transformation by returning:
593
/// - The resulting (possibly transformed) node,
594
/// - `transformed`: flag indicating whether any change was made to the node
595
/// - `tnr`: [`TreeNodeRecursion`] specifying how to proceed with the recursion.
596
///
597
/// At the end of the transformation, the return value will contain:
598
/// - The final (possibly transformed) tree,
599
/// - `transformed`: flag indicating whether any change was made to the node
600
/// - `tnr`: [`TreeNodeRecursion`] specifying how the recursion ended.
601
///
602
/// See also
603
/// * [`Transformed::update_data`] to modify the node without changing the `transformed` flag
604
/// * [`Transformed::map_data`] for fallable operation that return the same type
605
/// * [`Transformed::transform_data`] to chain fallable transformations
606
/// * [`TransformedResult`] for working with `Result<Transformed<U>>`
607
///
608
/// # Examples
609
///
610
/// Use [`Transformed::yes`] and [`Transformed::no`] to signal that a node was
611
/// rewritten and the recursion should continue:
612
///
613
/// ```
614
/// # use datafusion_common::tree_node::Transformed;
615
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
616
/// # fn orig_expr() -> i64 { 1 }
617
/// # fn make_new_expr(i: i64) -> i64 { 2 }
618
/// let expr = orig_expr();
619
///
620
/// // Create a new `Transformed` object signaling the node was not rewritten
621
/// let ret = Transformed::no(expr.clone());
622
/// assert!(!ret.transformed);
623
///
624
/// // Create a new `Transformed` object signaling the node was rewritten
625
/// let ret = Transformed::yes(expr);
626
/// assert!(ret.transformed)
627
/// ```
628
///
629
/// Access the node within the `Transformed` object:
630
/// ```
631
/// # use datafusion_common::tree_node::Transformed;
632
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
633
/// # fn orig_expr() -> i64 { 1 }
634
/// # fn make_new_expr(i: i64) -> i64 { 2 }
635
/// let expr = orig_expr();
636
///
637
/// // `Transformed` object signaling the node was not rewritten
638
/// let ret = Transformed::no(expr.clone());
639
/// // Access the inner object using .data
640
/// assert_eq!(expr, ret.data);
641
/// ```
642
///
643
/// Transform the node within the `Transformed` object.
644
///
645
/// ```
646
/// # use datafusion_common::tree_node::Transformed;
647
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
648
/// # fn orig_expr() -> i64 { 1 }
649
/// # fn make_new_expr(i: i64) -> i64 { 2 }
650
/// let expr = orig_expr();
651
/// let ret = Transformed::no(expr.clone())
652
///   .transform_data(|expr| {
653
///    // closure returns a result and potentially transforms the node
654
///    // in this example, it does transform the node
655
///    let new_expr = make_new_expr(expr);
656
///    Ok(Transformed::yes(new_expr))
657
///  }).unwrap();
658
/// // transformed flag is the union of the original ans closure's  transformed flag
659
/// assert!(ret.transformed);
660
/// ```
661
/// # Example APIs that use `TreeNode`
662
/// - [`TreeNode`],
663
/// - [`TreeNode::rewrite`],
664
/// - [`TreeNode::transform_down`],
665
/// - [`TreeNode::transform_up`],
666
/// - [`TreeNode::transform_down_up`]
667
#[derive(PartialEq, Debug)]
668
pub struct Transformed<T> {
669
    pub data: T,
670
    pub transformed: bool,
671
    pub tnr: TreeNodeRecursion,
672
}
673
674
impl<T> Transformed<T> {
675
    /// Create a new `Transformed` object with the given information.
676
61.9k
    pub fn new(data: T, transformed: bool, tnr: TreeNodeRecursion) -> Self {
677
61.9k
        Self {
678
61.9k
            data,
679
61.9k
            transformed,
680
61.9k
            tnr,
681
61.9k
        }
682
61.9k
    }
683
684
    /// Create a `Transformed` with `transformed` and [`TreeNodeRecursion::Continue`].
685
0
    pub fn new_transformed(data: T, transformed: bool) -> Self {
686
0
        Self::new(data, transformed, TreeNodeRecursion::Continue)
687
0
    }
688
689
    /// Wrapper for transformed data with [`TreeNodeRecursion::Continue`] statement.
690
20.4k
    pub fn yes(data: T) -> Self {
691
20.4k
        Self::new(data, true, TreeNodeRecursion::Continue)
692
20.4k
    }
693
694
    /// Wrapper for unchanged data with [`TreeNodeRecursion::Continue`] statement.
695
25.1k
    pub fn no(data: T) -> Self {
696
25.1k
        Self::new(data, false, TreeNodeRecursion::Continue)
697
25.1k
    }
698
699
    /// Applies an infallible `f` to the data of this [`Transformed`] object,
700
    /// without modifying the `transformed` flag.
701
0
    pub fn update_data<U, F: FnOnce(T) -> U>(self, f: F) -> Transformed<U> {
702
0
        Transformed::new(f(self.data), self.transformed, self.tnr)
703
0
    }
704
705
    /// Applies a fallible `f` (returns `Result`) to the data of this
706
    /// [`Transformed`] object, without modifying the `transformed` flag.
707
8.17k
    pub fn map_data<U, F: FnOnce(T) -> Result<U>>(self, f: F) -> Result<Transformed<U>> {
708
8.17k
        f(self.data).map(|data| Transformed::new(data, self.transformed, self.tnr))
709
8.17k
    }
710
711
    /// Applies a fallible transforming `f` to the data of this [`Transformed`]
712
    /// object.
713
    ///
714
    /// The returned `Transformed` object has the `transformed` flag set if either
715
    /// `self` or the return value of `f` have the `transformed` flag set.
716
0
    pub fn transform_data<U, F: FnOnce(T) -> Result<Transformed<U>>>(
717
0
        self,
718
0
        f: F,
719
0
    ) -> Result<Transformed<U>> {
720
0
        f(self.data).map(|mut t| {
721
0
            t.transformed |= self.transformed;
722
0
            t
723
0
        })
724
0
    }
725
726
    /// Maps the [`Transformed`] object to the result of the given `f` depending on the
727
    /// current [`TreeNodeRecursion`] value and the fact that `f` is changing the current
728
    /// node's children.
729
363
    pub fn transform_children<F: FnOnce(T) -> Result<Transformed<T>>>(
730
363
        mut self,
731
363
        f: F,
732
363
    ) -> Result<Transformed<T>> {
733
363
        match self.tnr {
734
            TreeNodeRecursion::Continue => {
735
363
                return f(self.data).map(|mut t| {
736
363
                    t.transformed |= self.transformed;
737
363
                    t
738
363
                });
739
            }
740
0
            TreeNodeRecursion::Jump => {
741
0
                self.tnr = TreeNodeRecursion::Continue;
742
0
            }
743
0
            TreeNodeRecursion::Stop => {}
744
        }
745
0
        Ok(self)
746
363
    }
747
748
    /// Maps the [`Transformed`] object to the result of the given `f` depending on the
749
    /// current [`TreeNodeRecursion`] value and the fact that `f` is changing the current
750
    /// node's sibling.
751
0
    pub fn transform_sibling<F: FnOnce(T) -> Result<Transformed<T>>>(
752
0
        self,
753
0
        f: F,
754
0
    ) -> Result<Transformed<T>> {
755
0
        match self.tnr {
756
            TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
757
0
                f(self.data).map(|mut t| {
758
0
                    t.transformed |= self.transformed;
759
0
                    t
760
0
                })
761
            }
762
0
            TreeNodeRecursion::Stop => Ok(self),
763
        }
764
0
    }
765
766
    /// Maps the [`Transformed`] object to the result of the given `f` depending on the
767
    /// current [`TreeNodeRecursion`] value and the fact that `f` is changing the current
768
    /// node's parent.
769
26.5k
    pub fn transform_parent<F: FnOnce(T) -> Result<Transformed<T>>>(
770
26.5k
        self,
771
26.5k
        f: F,
772
26.5k
    ) -> Result<Transformed<T>> {
773
26.5k
        match self.tnr {
774
26.5k
            TreeNodeRecursion::Continue => f(self.data).map(|mut t| {
775
26.5k
                t.transformed |= self.transformed;
776
26.5k
                t
777
26.5k
            }),
778
0
            TreeNodeRecursion::Jump | TreeNodeRecursion::Stop => Ok(self),
779
        }
780
26.5k
    }
781
}
782
783
/// Transformation helper to process a sequence of iterable tree nodes that are siblings.
784
pub trait TreeNodeIterator: Iterator {
785
    /// Apples `f` to each item in this iterator
786
    ///
787
    /// Visits all items in the iterator unless
788
    /// `f` returns an error or `f` returns `TreeNodeRecursion::Stop`.
789
    ///
790
    /// # Returns
791
    /// Error if `f` returns an error or `Ok(TreeNodeRecursion)` from the last invocation
792
    /// of `f` or `Continue` if the iterator is empty
793
    fn apply_until_stop<F: FnMut(Self::Item) -> Result<TreeNodeRecursion>>(
794
        self,
795
        f: F,
796
    ) -> Result<TreeNodeRecursion>;
797
798
    /// Apples `f` to each item in this iterator
799
    ///
800
    /// Visits all items in the iterator unless
801
    /// `f` returns an error or `f` returns `TreeNodeRecursion::Stop`.
802
    ///
803
    /// # Returns
804
    /// Error if `f` returns an error
805
    ///
806
    /// Ok(Transformed) such that:
807
    /// 1. `transformed` is true if any return from `f` had transformed true
808
    /// 2. `data` from the last invocation of `f`
809
    /// 3. `tnr` from the last invocation of `f` or `Continue` if the iterator is empty
810
    fn map_until_stop_and_collect<
811
        F: FnMut(Self::Item) -> Result<Transformed<Self::Item>>,
812
    >(
813
        self,
814
        f: F,
815
    ) -> Result<Transformed<Vec<Self::Item>>>;
816
}
817
818
impl<I: Iterator> TreeNodeIterator for I {
819
4.90k
    fn apply_until_stop<F: FnMut(Self::Item) -> Result<TreeNodeRecursion>>(
820
4.90k
        self,
821
4.90k
        mut f: F,
822
4.90k
    ) -> Result<TreeNodeRecursion> {
823
4.90k
        let mut tnr = TreeNodeRecursion::Continue;
824
5.30k
        for 
i402
in self {
825
402
            tnr = f(i)
?0
;
826
402
            match tnr {
827
402
                TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {}
828
0
                TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop),
829
            }
830
        }
831
4.90k
        Ok(tnr)
832
4.90k
    }
833
834
8.17k
    fn map_until_stop_and_collect<
835
8.17k
        F: FnMut(Self::Item) -> Result<Transformed<Self::Item>>,
836
8.17k
    >(
837
8.17k
        self,
838
8.17k
        mut f: F,
839
8.17k
    ) -> Result<Transformed<Vec<Self::Item>>> {
840
8.17k
        let mut tnr = TreeNodeRecursion::Continue;
841
8.17k
        let mut transformed = false;
842
16.0k
        self.map(|item| match tnr {
843
            TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {
844
16.0k
                f(item).map(|result| {
845
16.0k
                    tnr = result.tnr;
846
16.0k
                    transformed |= result.transformed;
847
16.0k
                    result.data
848
16.0k
                })
849
            }
850
0
            TreeNodeRecursion::Stop => Ok(item),
851
16.0k
        })
852
8.17k
        .collect::<Result<Vec<_>>>()
853
8.17k
        .map(|data| Transformed::new(data, transformed, tnr))
854
8.17k
    }
855
}
856
857
/// Transformation helper to process a heterogeneous sequence of tree node containing
858
/// expressions.
859
///
860
/// This macro is very similar to [TreeNodeIterator::map_until_stop_and_collect] to
861
/// process nodes that are siblings, but it accepts an initial transformation (`F0`) and
862
/// a sequence of pairs. Each pair is made of an expression (`EXPR`) and its
863
/// transformation (`F`).
864
///
865
/// The macro builds up a tuple that contains `Transformed.data` result of `F0` as the
866
/// first element and further elements from the sequence of pairs. An element from a pair
867
/// is either the value of `EXPR` or the `Transformed.data` result of `F`, depending on
868
/// the `Transformed.tnr` result of previous `F`s (`F0` initially).
869
///
870
/// # Returns
871
/// Error if any of the transformations returns an error
872
///
873
/// Ok(Transformed<(data0, ..., dataN)>) such that:
874
/// 1. `transformed` is true if any of the transformations had transformed true
875
/// 2. `(data0, ..., dataN)`, where `data0` is the `Transformed.data` from `F0` and
876
///     `data1` ... `dataN` are from either `EXPR` or the `Transformed.data` of `F`
877
/// 3. `tnr` from `F0` or the last invocation of `F`
878
#[macro_export]
879
macro_rules! map_until_stop_and_collect {
880
    ($F0:expr, $($EXPR:expr, $F:expr),*) => {{
881
0
        $F0.and_then(|Transformed { data: data0, mut transformed, mut tnr }| {
882
0
            let all_datas = (
883
0
                data0,
884
0
                $(
885
0
                    if tnr == TreeNodeRecursion::Continue || tnr == TreeNodeRecursion::Jump {
886
0
                        $F.map(|result| {
887
0
                            tnr = result.tnr;
888
0
                            transformed |= result.transformed;
889
0
                            result.data
890
0
                        })?
891
                    } else {
892
                        $EXPR
893
                    },
894
                )*
895
            );
896
0
            Ok(Transformed::new(all_datas, transformed, tnr))
897
0
        })
898
0
    }}
899
0
}
900
0
901
0
/// Transformation helper to access [`Transformed`] fields in a [`Result`] easily.
902
0
///
903
0
/// # Example
904
0
/// Access the internal data of a `Result<Transformed<T>>`
905
0
/// as a `Result<T>` using the `data` method:
906
0
/// ```
907
0
/// # use datafusion_common::Result;
908
0
/// # use datafusion_common::tree_node::{Transformed, TransformedResult};
909
0
/// # // note use i64 instead of Expr as Expr is not in datafusion-common
910
0
/// # fn update_expr() -> i64 { 1 }
911
0
/// # fn main() -> Result<()> {
912
0
/// let transformed: Result<Transformed<_>> = Ok(Transformed::yes(update_expr()));
913
0
/// // access the internal data of the transformed result, or return the error
914
0
/// let transformed_expr = transformed.data()?;
915
0
/// # Ok(())
916
0
/// # }
917
0
/// ```
918
0
pub trait TransformedResult<T> {
919
0
    fn data(self) -> Result<T>;
920
0
921
0
    fn transformed(self) -> Result<bool>;
922
0
923
0
    fn tnr(self) -> Result<TreeNodeRecursion>;
924
0
}
925
0
926
0
impl<T> TransformedResult<T> for Result<Transformed<T>> {
927
10.8k
    fn data(self) -> Result<T> {
928
10.8k
        self.map(|t| t.data)
929
10.8k
    }
930
931
0
    fn transformed(self) -> Result<bool> {
932
0
        self.map(|t| t.transformed)
933
0
    }
934
935
0
    fn tnr(self) -> Result<TreeNodeRecursion> {
936
0
        self.map(|t| t.tnr)
937
0
    }
938
}
939
940
/// Helper trait for implementing [`TreeNode`] that have children stored as
941
/// `Arc`s. If some trait object, such as `dyn T`, implements this trait,
942
/// its related `Arc<dyn T>` will automatically implement [`TreeNode`].
943
pub trait DynTreeNode {
944
    /// Returns all children of the specified `TreeNode`.
945
    fn arc_children(&self) -> Vec<&Arc<Self>>;
946
947
    /// Constructs a new node with the specified children.
948
    fn with_new_arc_children(
949
        &self,
950
        arc_self: Arc<Self>,
951
        new_children: Vec<Arc<Self>>,
952
    ) -> Result<Arc<Self>>;
953
}
954
955
/// Blanket implementation for any `Arc<T>` where `T` implements [`DynTreeNode`]
956
/// (such as [`Arc<dyn PhysicalExpr>`]).
957
impl<T: DynTreeNode + ?Sized> TreeNode for Arc<T> {
958
4.90k
    fn apply_children<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
959
4.90k
        &'n self,
960
4.90k
        f: F,
961
4.90k
    ) -> Result<TreeNodeRecursion> {
962
4.90k
        self.arc_children().into_iter().apply_until_stop(f)
963
4.90k
    }
964
965
9.34k
    fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>(
966
9.34k
        self,
967
9.34k
        f: F,
968
9.34k
    ) -> Result<Transformed<Self>> {
969
9.34k
        let children = self.arc_children();
970
9.34k
        if !children.is_empty() {
971
66
            let new_children = children
972
66
                .into_iter()
973
66
                .cloned()
974
66
                .map_until_stop_and_collect(f)
?0
;
975
            // Propagate up `new_children.transformed` and `new_children.tnr`
976
            // along with the node containing transformed children.
977
66
            if new_children.transformed {
978
66
                let arc_self = Arc::clone(&self);
979
66
                new_children.map_data(|new_children| {
980
66
                    self.with_new_arc_children(arc_self, new_children)
981
66
                })
982
            } else {
983
0
                Ok(Transformed::new(self, false, new_children.tnr))
984
            }
985
        } else {
986
9.27k
            Ok(Transformed::no(self))
987
        }
988
9.34k
    }
989
}
990
991
/// Instead of implementing [`TreeNode`], it's recommended to implement a [`ConcreteTreeNode`] for
992
/// trees that contain nodes with payloads. This approach ensures safe execution of algorithms
993
/// involving payloads, by enforcing rules for detaching and reattaching child nodes.
994
pub trait ConcreteTreeNode: Sized {
995
    /// Provides read-only access to child nodes.
996
    fn children(&self) -> &[Self];
997
998
    /// Detaches the node from its children, returning the node itself and its detached children.
999
    fn take_children(self) -> (Self, Vec<Self>);
1000
1001
    /// Reattaches updated child nodes to the node, returning the updated node.
1002
    fn with_new_children(self, children: Vec<Self>) -> Result<Self>;
1003
}
1004
1005
impl<T: ConcreteTreeNode> TreeNode for T {
1006
0
    fn apply_children<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
1007
0
        &'n self,
1008
0
        f: F,
1009
0
    ) -> Result<TreeNodeRecursion> {
1010
0
        self.children().iter().apply_until_stop(f)
1011
0
    }
1012
1013
17.5k
    fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>(
1014
17.5k
        self,
1015
17.5k
        f: F,
1016
17.5k
    ) -> Result<Transformed<Self>> {
1017
17.5k
        let (new_self, children) = self.take_children();
1018
17.5k
        if !children.is_empty() {
1019
8.10k
            let new_children = children.into_iter().map_until_stop_and_collect(f)
?0
;
1020
            // Propagate up `new_children.transformed` and `new_children.tnr` along with
1021
            // the node containing transformed children.
1022
8.10k
            new_children.map_data(|new_children| new_self.with_new_children(new_children))
1023
        } else {
1024
9.42k
            Ok(Transformed::no(new_self))
1025
        }
1026
17.5k
    }
1027
}
1028
1029
#[cfg(test)]
1030
mod tests {
1031
    use std::collections::HashMap;
1032
    use std::fmt::Display;
1033
1034
    use crate::tree_node::{
1035
        Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion, TreeNodeRewriter,
1036
        TreeNodeVisitor,
1037
    };
1038
    use crate::Result;
1039
1040
    #[derive(Debug, Eq, Hash, PartialEq)]
1041
    struct TestTreeNode<T> {
1042
        children: Vec<TestTreeNode<T>>,
1043
        data: T,
1044
    }
1045
1046
    impl<T> TestTreeNode<T> {
1047
        fn new(children: Vec<TestTreeNode<T>>, data: T) -> Self {
1048
            Self { children, data }
1049
        }
1050
    }
1051
1052
    impl<T> TreeNode for TestTreeNode<T> {
1053
        fn apply_children<'n, F: FnMut(&'n Self) -> Result<TreeNodeRecursion>>(
1054
            &'n self,
1055
            f: F,
1056
        ) -> Result<TreeNodeRecursion> {
1057
            self.children.iter().apply_until_stop(f)
1058
        }
1059
1060
        fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>(
1061
            self,
1062
            f: F,
1063
        ) -> Result<Transformed<Self>> {
1064
            Ok(self
1065
                .children
1066
                .into_iter()
1067
                .map_until_stop_and_collect(f)?
1068
                .update_data(|new_children| Self {
1069
                    children: new_children,
1070
                    ..self
1071
                }))
1072
        }
1073
    }
1074
1075
    //       J
1076
    //       |
1077
    //       I
1078
    //       |
1079
    //       F
1080
    //     /   \
1081
    //    E     G
1082
    //    |     |
1083
    //    C     H
1084
    //  /   \
1085
    // B     D
1086
    //       |
1087
    //       A
1088
    fn test_tree() -> TestTreeNode<String> {
1089
        let node_a = TestTreeNode::new(vec![], "a".to_string());
1090
        let node_b = TestTreeNode::new(vec![], "b".to_string());
1091
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
1092
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
1093
        let node_e = TestTreeNode::new(vec![node_c], "e".to_string());
1094
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1095
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1096
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f".to_string());
1097
        let node_i = TestTreeNode::new(vec![node_f], "i".to_string());
1098
        TestTreeNode::new(vec![node_i], "j".to_string())
1099
    }
1100
1101
    // Continue on all nodes
1102
    // Expected visits in a combined traversal
1103
    fn all_visits() -> Vec<String> {
1104
        vec![
1105
            "f_down(j)",
1106
            "f_down(i)",
1107
            "f_down(f)",
1108
            "f_down(e)",
1109
            "f_down(c)",
1110
            "f_down(b)",
1111
            "f_up(b)",
1112
            "f_down(d)",
1113
            "f_down(a)",
1114
            "f_up(a)",
1115
            "f_up(d)",
1116
            "f_up(c)",
1117
            "f_up(e)",
1118
            "f_down(g)",
1119
            "f_down(h)",
1120
            "f_up(h)",
1121
            "f_up(g)",
1122
            "f_up(f)",
1123
            "f_up(i)",
1124
            "f_up(j)",
1125
        ]
1126
        .into_iter()
1127
        .map(|s| s.to_string())
1128
        .collect()
1129
    }
1130
1131
    // Expected transformed tree after a combined traversal
1132
    fn transformed_tree() -> TestTreeNode<String> {
1133
        let node_a = TestTreeNode::new(vec![], "f_up(f_down(a))".to_string());
1134
        let node_b = TestTreeNode::new(vec![], "f_up(f_down(b))".to_string());
1135
        let node_d = TestTreeNode::new(vec![node_a], "f_up(f_down(d))".to_string());
1136
        let node_c =
1137
            TestTreeNode::new(vec![node_b, node_d], "f_up(f_down(c))".to_string());
1138
        let node_e = TestTreeNode::new(vec![node_c], "f_up(f_down(e))".to_string());
1139
        let node_h = TestTreeNode::new(vec![], "f_up(f_down(h))".to_string());
1140
        let node_g = TestTreeNode::new(vec![node_h], "f_up(f_down(g))".to_string());
1141
        let node_f =
1142
            TestTreeNode::new(vec![node_e, node_g], "f_up(f_down(f))".to_string());
1143
        let node_i = TestTreeNode::new(vec![node_f], "f_up(f_down(i))".to_string());
1144
        TestTreeNode::new(vec![node_i], "f_up(f_down(j))".to_string())
1145
    }
1146
1147
    // Expected transformed tree after a top-down traversal
1148
    fn transformed_down_tree() -> TestTreeNode<String> {
1149
        let node_a = TestTreeNode::new(vec![], "f_down(a)".to_string());
1150
        let node_b = TestTreeNode::new(vec![], "f_down(b)".to_string());
1151
        let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string());
1152
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string());
1153
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1154
        let node_h = TestTreeNode::new(vec![], "f_down(h)".to_string());
1155
        let node_g = TestTreeNode::new(vec![node_h], "f_down(g)".to_string());
1156
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1157
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1158
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1159
    }
1160
1161
    // Expected transformed tree after a bottom-up traversal
1162
    fn transformed_up_tree() -> TestTreeNode<String> {
1163
        let node_a = TestTreeNode::new(vec![], "f_up(a)".to_string());
1164
        let node_b = TestTreeNode::new(vec![], "f_up(b)".to_string());
1165
        let node_d = TestTreeNode::new(vec![node_a], "f_up(d)".to_string());
1166
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_up(c)".to_string());
1167
        let node_e = TestTreeNode::new(vec![node_c], "f_up(e)".to_string());
1168
        let node_h = TestTreeNode::new(vec![], "f_up(h)".to_string());
1169
        let node_g = TestTreeNode::new(vec![node_h], "f_up(g)".to_string());
1170
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_up(f)".to_string());
1171
        let node_i = TestTreeNode::new(vec![node_f], "f_up(i)".to_string());
1172
        TestTreeNode::new(vec![node_i], "f_up(j)".to_string())
1173
    }
1174
1175
    // f_down Jump on A node
1176
    fn f_down_jump_on_a_visits() -> Vec<String> {
1177
        vec![
1178
            "f_down(j)",
1179
            "f_down(i)",
1180
            "f_down(f)",
1181
            "f_down(e)",
1182
            "f_down(c)",
1183
            "f_down(b)",
1184
            "f_up(b)",
1185
            "f_down(d)",
1186
            "f_down(a)",
1187
            "f_up(a)",
1188
            "f_up(d)",
1189
            "f_up(c)",
1190
            "f_up(e)",
1191
            "f_down(g)",
1192
            "f_down(h)",
1193
            "f_up(h)",
1194
            "f_up(g)",
1195
            "f_up(f)",
1196
            "f_up(i)",
1197
            "f_up(j)",
1198
        ]
1199
        .into_iter()
1200
        .map(|s| s.to_string())
1201
        .collect()
1202
    }
1203
1204
    fn f_down_jump_on_a_transformed_down_tree() -> TestTreeNode<String> {
1205
        let node_a = TestTreeNode::new(vec![], "f_down(a)".to_string());
1206
        let node_b = TestTreeNode::new(vec![], "f_down(b)".to_string());
1207
        let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string());
1208
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string());
1209
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1210
        let node_h = TestTreeNode::new(vec![], "f_down(h)".to_string());
1211
        let node_g = TestTreeNode::new(vec![node_h], "f_down(g)".to_string());
1212
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1213
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1214
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1215
    }
1216
1217
    // f_down Jump on E node
1218
    fn f_down_jump_on_e_visits() -> Vec<String> {
1219
        vec![
1220
            "f_down(j)",
1221
            "f_down(i)",
1222
            "f_down(f)",
1223
            "f_down(e)",
1224
            "f_up(e)",
1225
            "f_down(g)",
1226
            "f_down(h)",
1227
            "f_up(h)",
1228
            "f_up(g)",
1229
            "f_up(f)",
1230
            "f_up(i)",
1231
            "f_up(j)",
1232
        ]
1233
        .into_iter()
1234
        .map(|s| s.to_string())
1235
        .collect()
1236
    }
1237
1238
    fn f_down_jump_on_e_transformed_tree() -> TestTreeNode<String> {
1239
        let node_a = TestTreeNode::new(vec![], "a".to_string());
1240
        let node_b = TestTreeNode::new(vec![], "b".to_string());
1241
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
1242
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
1243
        let node_e = TestTreeNode::new(vec![node_c], "f_up(f_down(e))".to_string());
1244
        let node_h = TestTreeNode::new(vec![], "f_up(f_down(h))".to_string());
1245
        let node_g = TestTreeNode::new(vec![node_h], "f_up(f_down(g))".to_string());
1246
        let node_f =
1247
            TestTreeNode::new(vec![node_e, node_g], "f_up(f_down(f))".to_string());
1248
        let node_i = TestTreeNode::new(vec![node_f], "f_up(f_down(i))".to_string());
1249
        TestTreeNode::new(vec![node_i], "f_up(f_down(j))".to_string())
1250
    }
1251
1252
    fn f_down_jump_on_e_transformed_down_tree() -> TestTreeNode<String> {
1253
        let node_a = TestTreeNode::new(vec![], "a".to_string());
1254
        let node_b = TestTreeNode::new(vec![], "b".to_string());
1255
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
1256
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
1257
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1258
        let node_h = TestTreeNode::new(vec![], "f_down(h)".to_string());
1259
        let node_g = TestTreeNode::new(vec![node_h], "f_down(g)".to_string());
1260
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1261
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1262
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1263
    }
1264
1265
    // f_up Jump on A node
1266
    fn f_up_jump_on_a_visits() -> Vec<String> {
1267
        vec![
1268
            "f_down(j)",
1269
            "f_down(i)",
1270
            "f_down(f)",
1271
            "f_down(e)",
1272
            "f_down(c)",
1273
            "f_down(b)",
1274
            "f_up(b)",
1275
            "f_down(d)",
1276
            "f_down(a)",
1277
            "f_up(a)",
1278
            "f_down(g)",
1279
            "f_down(h)",
1280
            "f_up(h)",
1281
            "f_up(g)",
1282
            "f_up(f)",
1283
            "f_up(i)",
1284
            "f_up(j)",
1285
        ]
1286
        .into_iter()
1287
        .map(|s| s.to_string())
1288
        .collect()
1289
    }
1290
1291
    fn f_up_jump_on_a_transformed_tree() -> TestTreeNode<String> {
1292
        let node_a = TestTreeNode::new(vec![], "f_up(f_down(a))".to_string());
1293
        let node_b = TestTreeNode::new(vec![], "f_up(f_down(b))".to_string());
1294
        let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string());
1295
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string());
1296
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1297
        let node_h = TestTreeNode::new(vec![], "f_up(f_down(h))".to_string());
1298
        let node_g = TestTreeNode::new(vec![node_h], "f_up(f_down(g))".to_string());
1299
        let node_f =
1300
            TestTreeNode::new(vec![node_e, node_g], "f_up(f_down(f))".to_string());
1301
        let node_i = TestTreeNode::new(vec![node_f], "f_up(f_down(i))".to_string());
1302
        TestTreeNode::new(vec![node_i], "f_up(f_down(j))".to_string())
1303
    }
1304
1305
    fn f_up_jump_on_a_transformed_up_tree() -> TestTreeNode<String> {
1306
        let node_a = TestTreeNode::new(vec![], "f_up(a)".to_string());
1307
        let node_b = TestTreeNode::new(vec![], "f_up(b)".to_string());
1308
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
1309
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
1310
        let node_e = TestTreeNode::new(vec![node_c], "e".to_string());
1311
        let node_h = TestTreeNode::new(vec![], "f_up(h)".to_string());
1312
        let node_g = TestTreeNode::new(vec![node_h], "f_up(g)".to_string());
1313
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_up(f)".to_string());
1314
        let node_i = TestTreeNode::new(vec![node_f], "f_up(i)".to_string());
1315
        TestTreeNode::new(vec![node_i], "f_up(j)".to_string())
1316
    }
1317
1318
    // f_up Jump on E node
1319
    fn f_up_jump_on_e_visits() -> Vec<String> {
1320
        vec![
1321
            "f_down(j)",
1322
            "f_down(i)",
1323
            "f_down(f)",
1324
            "f_down(e)",
1325
            "f_down(c)",
1326
            "f_down(b)",
1327
            "f_up(b)",
1328
            "f_down(d)",
1329
            "f_down(a)",
1330
            "f_up(a)",
1331
            "f_up(d)",
1332
            "f_up(c)",
1333
            "f_up(e)",
1334
            "f_down(g)",
1335
            "f_down(h)",
1336
            "f_up(h)",
1337
            "f_up(g)",
1338
            "f_up(f)",
1339
            "f_up(i)",
1340
            "f_up(j)",
1341
        ]
1342
        .into_iter()
1343
        .map(|s| s.to_string())
1344
        .collect()
1345
    }
1346
1347
    fn f_up_jump_on_e_transformed_tree() -> TestTreeNode<String> {
1348
        transformed_tree()
1349
    }
1350
1351
    fn f_up_jump_on_e_transformed_up_tree() -> TestTreeNode<String> {
1352
        transformed_up_tree()
1353
    }
1354
1355
    // f_down Stop on A node
1356
1357
    fn f_down_stop_on_a_visits() -> Vec<String> {
1358
        vec![
1359
            "f_down(j)",
1360
            "f_down(i)",
1361
            "f_down(f)",
1362
            "f_down(e)",
1363
            "f_down(c)",
1364
            "f_down(b)",
1365
            "f_up(b)",
1366
            "f_down(d)",
1367
            "f_down(a)",
1368
        ]
1369
        .into_iter()
1370
        .map(|s| s.to_string())
1371
        .collect()
1372
    }
1373
1374
    fn f_down_stop_on_a_transformed_tree() -> TestTreeNode<String> {
1375
        let node_a = TestTreeNode::new(vec![], "f_down(a)".to_string());
1376
        let node_b = TestTreeNode::new(vec![], "f_up(f_down(b))".to_string());
1377
        let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string());
1378
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string());
1379
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1380
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1381
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1382
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1383
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1384
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1385
    }
1386
1387
    fn f_down_stop_on_a_transformed_down_tree() -> TestTreeNode<String> {
1388
        let node_a = TestTreeNode::new(vec![], "f_down(a)".to_string());
1389
        let node_b = TestTreeNode::new(vec![], "f_down(b)".to_string());
1390
        let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string());
1391
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string());
1392
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1393
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1394
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1395
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1396
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1397
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1398
    }
1399
1400
    // f_down Stop on E node
1401
    fn f_down_stop_on_e_visits() -> Vec<String> {
1402
        vec!["f_down(j)", "f_down(i)", "f_down(f)", "f_down(e)"]
1403
            .into_iter()
1404
            .map(|s| s.to_string())
1405
            .collect()
1406
    }
1407
1408
    fn f_down_stop_on_e_transformed_tree() -> TestTreeNode<String> {
1409
        let node_a = TestTreeNode::new(vec![], "a".to_string());
1410
        let node_b = TestTreeNode::new(vec![], "b".to_string());
1411
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
1412
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
1413
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1414
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1415
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1416
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1417
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1418
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1419
    }
1420
1421
    fn f_down_stop_on_e_transformed_down_tree() -> TestTreeNode<String> {
1422
        let node_a = TestTreeNode::new(vec![], "a".to_string());
1423
        let node_b = TestTreeNode::new(vec![], "b".to_string());
1424
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
1425
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
1426
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1427
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1428
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1429
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1430
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1431
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1432
    }
1433
1434
    // f_up Stop on A node
1435
    fn f_up_stop_on_a_visits() -> Vec<String> {
1436
        vec![
1437
            "f_down(j)",
1438
            "f_down(i)",
1439
            "f_down(f)",
1440
            "f_down(e)",
1441
            "f_down(c)",
1442
            "f_down(b)",
1443
            "f_up(b)",
1444
            "f_down(d)",
1445
            "f_down(a)",
1446
            "f_up(a)",
1447
        ]
1448
        .into_iter()
1449
        .map(|s| s.to_string())
1450
        .collect()
1451
    }
1452
1453
    fn f_up_stop_on_a_transformed_tree() -> TestTreeNode<String> {
1454
        let node_a = TestTreeNode::new(vec![], "f_up(f_down(a))".to_string());
1455
        let node_b = TestTreeNode::new(vec![], "f_up(f_down(b))".to_string());
1456
        let node_d = TestTreeNode::new(vec![node_a], "f_down(d)".to_string());
1457
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_down(c)".to_string());
1458
        let node_e = TestTreeNode::new(vec![node_c], "f_down(e)".to_string());
1459
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1460
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1461
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1462
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1463
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1464
    }
1465
1466
    fn f_up_stop_on_a_transformed_up_tree() -> TestTreeNode<String> {
1467
        let node_a = TestTreeNode::new(vec![], "f_up(a)".to_string());
1468
        let node_b = TestTreeNode::new(vec![], "f_up(b)".to_string());
1469
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
1470
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
1471
        let node_e = TestTreeNode::new(vec![node_c], "e".to_string());
1472
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1473
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1474
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f".to_string());
1475
        let node_i = TestTreeNode::new(vec![node_f], "i".to_string());
1476
        TestTreeNode::new(vec![node_i], "j".to_string())
1477
    }
1478
1479
    // f_up Stop on E node
1480
    fn f_up_stop_on_e_visits() -> Vec<String> {
1481
        vec![
1482
            "f_down(j)",
1483
            "f_down(i)",
1484
            "f_down(f)",
1485
            "f_down(e)",
1486
            "f_down(c)",
1487
            "f_down(b)",
1488
            "f_up(b)",
1489
            "f_down(d)",
1490
            "f_down(a)",
1491
            "f_up(a)",
1492
            "f_up(d)",
1493
            "f_up(c)",
1494
            "f_up(e)",
1495
        ]
1496
        .into_iter()
1497
        .map(|s| s.to_string())
1498
        .collect()
1499
    }
1500
1501
    fn f_up_stop_on_e_transformed_tree() -> TestTreeNode<String> {
1502
        let node_a = TestTreeNode::new(vec![], "f_up(f_down(a))".to_string());
1503
        let node_b = TestTreeNode::new(vec![], "f_up(f_down(b))".to_string());
1504
        let node_d = TestTreeNode::new(vec![node_a], "f_up(f_down(d))".to_string());
1505
        let node_c =
1506
            TestTreeNode::new(vec![node_b, node_d], "f_up(f_down(c))".to_string());
1507
        let node_e = TestTreeNode::new(vec![node_c], "f_up(f_down(e))".to_string());
1508
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1509
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1510
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f_down(f)".to_string());
1511
        let node_i = TestTreeNode::new(vec![node_f], "f_down(i)".to_string());
1512
        TestTreeNode::new(vec![node_i], "f_down(j)".to_string())
1513
    }
1514
1515
    fn f_up_stop_on_e_transformed_up_tree() -> TestTreeNode<String> {
1516
        let node_a = TestTreeNode::new(vec![], "f_up(a)".to_string());
1517
        let node_b = TestTreeNode::new(vec![], "f_up(b)".to_string());
1518
        let node_d = TestTreeNode::new(vec![node_a], "f_up(d)".to_string());
1519
        let node_c = TestTreeNode::new(vec![node_b, node_d], "f_up(c)".to_string());
1520
        let node_e = TestTreeNode::new(vec![node_c], "f_up(e)".to_string());
1521
        let node_h = TestTreeNode::new(vec![], "h".to_string());
1522
        let node_g = TestTreeNode::new(vec![node_h], "g".to_string());
1523
        let node_f = TestTreeNode::new(vec![node_e, node_g], "f".to_string());
1524
        let node_i = TestTreeNode::new(vec![node_f], "i".to_string());
1525
        TestTreeNode::new(vec![node_i], "j".to_string())
1526
    }
1527
1528
    fn down_visits(visits: Vec<String>) -> Vec<String> {
1529
        visits
1530
            .into_iter()
1531
            .filter(|v| v.starts_with("f_down"))
1532
            .collect()
1533
    }
1534
1535
    type TestVisitorF<T> = Box<dyn FnMut(&TestTreeNode<T>) -> Result<TreeNodeRecursion>>;
1536
1537
    struct TestVisitor<T> {
1538
        visits: Vec<String>,
1539
        f_down: TestVisitorF<T>,
1540
        f_up: TestVisitorF<T>,
1541
    }
1542
1543
    impl<T> TestVisitor<T> {
1544
        fn new(f_down: TestVisitorF<T>, f_up: TestVisitorF<T>) -> Self {
1545
            Self {
1546
                visits: vec![],
1547
                f_down,
1548
                f_up,
1549
            }
1550
        }
1551
    }
1552
1553
    impl<'n, T: Display> TreeNodeVisitor<'n> for TestVisitor<T> {
1554
        type Node = TestTreeNode<T>;
1555
1556
        fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
1557
            self.visits.push(format!("f_down({})", node.data));
1558
            (*self.f_down)(node)
1559
        }
1560
1561
        fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
1562
            self.visits.push(format!("f_up({})", node.data));
1563
            (*self.f_up)(node)
1564
        }
1565
    }
1566
1567
    fn visit_continue<T>(_: &TestTreeNode<T>) -> Result<TreeNodeRecursion> {
1568
        Ok(TreeNodeRecursion::Continue)
1569
    }
1570
1571
    fn visit_event_on<T: PartialEq, D: Into<T>>(
1572
        data: D,
1573
        event: TreeNodeRecursion,
1574
    ) -> impl FnMut(&TestTreeNode<T>) -> Result<TreeNodeRecursion> {
1575
        let d = data.into();
1576
        move |node| {
1577
            Ok(if node.data == d {
1578
                event
1579
            } else {
1580
                TreeNodeRecursion::Continue
1581
            })
1582
        }
1583
    }
1584
1585
    macro_rules! visit_test {
1586
        ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_VISITS:expr) => {
1587
            #[test]
1588
            fn $NAME() -> Result<()> {
1589
                let tree = test_tree();
1590
                let mut visitor = TestVisitor::new(Box::new($F_DOWN), Box::new($F_UP));
1591
                tree.visit(&mut visitor)?;
1592
                assert_eq!(visitor.visits, $EXPECTED_VISITS);
1593
1594
                Ok(())
1595
            }
1596
        };
1597
    }
1598
1599
    macro_rules! test_apply {
1600
        ($NAME:ident, $F:expr, $EXPECTED_VISITS:expr) => {
1601
            #[test]
1602
            fn $NAME() -> Result<()> {
1603
                let tree = test_tree();
1604
                let mut visits = vec![];
1605
                tree.apply(|node| {
1606
                    visits.push(format!("f_down({})", node.data));
1607
                    $F(node)
1608
                })?;
1609
                assert_eq!(visits, $EXPECTED_VISITS);
1610
1611
                Ok(())
1612
            }
1613
        };
1614
    }
1615
1616
    type TestRewriterF<T> =
1617
        Box<dyn FnMut(TestTreeNode<T>) -> Result<Transformed<TestTreeNode<T>>>>;
1618
1619
    struct TestRewriter<T> {
1620
        f_down: TestRewriterF<T>,
1621
        f_up: TestRewriterF<T>,
1622
    }
1623
1624
    impl<T> TestRewriter<T> {
1625
        fn new(f_down: TestRewriterF<T>, f_up: TestRewriterF<T>) -> Self {
1626
            Self { f_down, f_up }
1627
        }
1628
    }
1629
1630
    impl<T: Display> TreeNodeRewriter for TestRewriter<T> {
1631
        type Node = TestTreeNode<T>;
1632
1633
        fn f_down(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
1634
            (*self.f_down)(node)
1635
        }
1636
1637
        fn f_up(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
1638
            (*self.f_up)(node)
1639
        }
1640
    }
1641
1642
    fn transform_yes<N: Display, T: Display + From<String>>(
1643
        transformation_name: N,
1644
    ) -> impl FnMut(TestTreeNode<T>) -> Result<Transformed<TestTreeNode<T>>> {
1645
        move |node| {
1646
            Ok(Transformed::yes(TestTreeNode::new(
1647
                node.children,
1648
                format!("{}({})", transformation_name, node.data).into(),
1649
            )))
1650
        }
1651
    }
1652
1653
    fn transform_and_event_on<
1654
        N: Display,
1655
        T: PartialEq + Display + From<String>,
1656
        D: Into<T>,
1657
    >(
1658
        transformation_name: N,
1659
        data: D,
1660
        event: TreeNodeRecursion,
1661
    ) -> impl FnMut(TestTreeNode<T>) -> Result<Transformed<TestTreeNode<T>>> {
1662
        let d = data.into();
1663
        move |node| {
1664
            let new_node = TestTreeNode::new(
1665
                node.children,
1666
                format!("{}({})", transformation_name, node.data).into(),
1667
            );
1668
            Ok(if node.data == d {
1669
                Transformed::new(new_node, true, event)
1670
            } else {
1671
                Transformed::yes(new_node)
1672
            })
1673
        }
1674
    }
1675
1676
    macro_rules! rewrite_test {
1677
        ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_TREE:expr) => {
1678
            #[test]
1679
            fn $NAME() -> Result<()> {
1680
                let tree = test_tree();
1681
                let mut rewriter = TestRewriter::new(Box::new($F_DOWN), Box::new($F_UP));
1682
                assert_eq!(tree.rewrite(&mut rewriter)?, $EXPECTED_TREE);
1683
1684
                Ok(())
1685
            }
1686
        };
1687
    }
1688
1689
    macro_rules! transform_test {
1690
        ($NAME:ident, $F_DOWN:expr, $F_UP:expr, $EXPECTED_TREE:expr) => {
1691
            #[test]
1692
            fn $NAME() -> Result<()> {
1693
                let tree = test_tree();
1694
                assert_eq!(tree.transform_down_up($F_DOWN, $F_UP,)?, $EXPECTED_TREE);
1695
1696
                Ok(())
1697
            }
1698
        };
1699
    }
1700
1701
    macro_rules! transform_down_test {
1702
        ($NAME:ident, $F:expr, $EXPECTED_TREE:expr) => {
1703
            #[test]
1704
            fn $NAME() -> Result<()> {
1705
                let tree = test_tree();
1706
                assert_eq!(tree.transform_down($F)?, $EXPECTED_TREE);
1707
1708
                Ok(())
1709
            }
1710
        };
1711
    }
1712
1713
    macro_rules! transform_up_test {
1714
        ($NAME:ident, $F:expr, $EXPECTED_TREE:expr) => {
1715
            #[test]
1716
            fn $NAME() -> Result<()> {
1717
                let tree = test_tree();
1718
                assert_eq!(tree.transform_up($F)?, $EXPECTED_TREE);
1719
1720
                Ok(())
1721
            }
1722
        };
1723
    }
1724
1725
    visit_test!(test_visit, visit_continue, visit_continue, all_visits());
1726
    visit_test!(
1727
        test_visit_f_down_jump_on_a,
1728
        visit_event_on("a", TreeNodeRecursion::Jump),
1729
        visit_continue,
1730
        f_down_jump_on_a_visits()
1731
    );
1732
    visit_test!(
1733
        test_visit_f_down_jump_on_e,
1734
        visit_event_on("e", TreeNodeRecursion::Jump),
1735
        visit_continue,
1736
        f_down_jump_on_e_visits()
1737
    );
1738
    visit_test!(
1739
        test_visit_f_up_jump_on_a,
1740
        visit_continue,
1741
        visit_event_on("a", TreeNodeRecursion::Jump),
1742
        f_up_jump_on_a_visits()
1743
    );
1744
    visit_test!(
1745
        test_visit_f_up_jump_on_e,
1746
        visit_continue,
1747
        visit_event_on("e", TreeNodeRecursion::Jump),
1748
        f_up_jump_on_e_visits()
1749
    );
1750
    visit_test!(
1751
        test_visit_f_down_stop_on_a,
1752
        visit_event_on("a", TreeNodeRecursion::Stop),
1753
        visit_continue,
1754
        f_down_stop_on_a_visits()
1755
    );
1756
    visit_test!(
1757
        test_visit_f_down_stop_on_e,
1758
        visit_event_on("e", TreeNodeRecursion::Stop),
1759
        visit_continue,
1760
        f_down_stop_on_e_visits()
1761
    );
1762
    visit_test!(
1763
        test_visit_f_up_stop_on_a,
1764
        visit_continue,
1765
        visit_event_on("a", TreeNodeRecursion::Stop),
1766
        f_up_stop_on_a_visits()
1767
    );
1768
    visit_test!(
1769
        test_visit_f_up_stop_on_e,
1770
        visit_continue,
1771
        visit_event_on("e", TreeNodeRecursion::Stop),
1772
        f_up_stop_on_e_visits()
1773
    );
1774
1775
    test_apply!(test_apply, visit_continue, down_visits(all_visits()));
1776
    test_apply!(
1777
        test_apply_f_down_jump_on_a,
1778
        visit_event_on("a", TreeNodeRecursion::Jump),
1779
        down_visits(f_down_jump_on_a_visits())
1780
    );
1781
    test_apply!(
1782
        test_apply_f_down_jump_on_e,
1783
        visit_event_on("e", TreeNodeRecursion::Jump),
1784
        down_visits(f_down_jump_on_e_visits())
1785
    );
1786
    test_apply!(
1787
        test_apply_f_down_stop_on_a,
1788
        visit_event_on("a", TreeNodeRecursion::Stop),
1789
        down_visits(f_down_stop_on_a_visits())
1790
    );
1791
    test_apply!(
1792
        test_apply_f_down_stop_on_e,
1793
        visit_event_on("e", TreeNodeRecursion::Stop),
1794
        down_visits(f_down_stop_on_e_visits())
1795
    );
1796
1797
    rewrite_test!(
1798
        test_rewrite,
1799
        transform_yes("f_down"),
1800
        transform_yes("f_up"),
1801
        Transformed::yes(transformed_tree())
1802
    );
1803
    rewrite_test!(
1804
        test_rewrite_f_down_jump_on_a,
1805
        transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump),
1806
        transform_yes("f_up"),
1807
        Transformed::yes(transformed_tree())
1808
    );
1809
    rewrite_test!(
1810
        test_rewrite_f_down_jump_on_e,
1811
        transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump),
1812
        transform_yes("f_up"),
1813
        Transformed::yes(f_down_jump_on_e_transformed_tree())
1814
    );
1815
    rewrite_test!(
1816
        test_rewrite_f_up_jump_on_a,
1817
        transform_yes("f_down"),
1818
        transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Jump),
1819
        Transformed::yes(f_up_jump_on_a_transformed_tree())
1820
    );
1821
    rewrite_test!(
1822
        test_rewrite_f_up_jump_on_e,
1823
        transform_yes("f_down"),
1824
        transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Jump),
1825
        Transformed::yes(f_up_jump_on_e_transformed_tree())
1826
    );
1827
    rewrite_test!(
1828
        test_rewrite_f_down_stop_on_a,
1829
        transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop),
1830
        transform_yes("f_up"),
1831
        Transformed::new(
1832
            f_down_stop_on_a_transformed_tree(),
1833
            true,
1834
            TreeNodeRecursion::Stop
1835
        )
1836
    );
1837
    rewrite_test!(
1838
        test_rewrite_f_down_stop_on_e,
1839
        transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop),
1840
        transform_yes("f_up"),
1841
        Transformed::new(
1842
            f_down_stop_on_e_transformed_tree(),
1843
            true,
1844
            TreeNodeRecursion::Stop
1845
        )
1846
    );
1847
    rewrite_test!(
1848
        test_rewrite_f_up_stop_on_a,
1849
        transform_yes("f_down"),
1850
        transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Stop),
1851
        Transformed::new(
1852
            f_up_stop_on_a_transformed_tree(),
1853
            true,
1854
            TreeNodeRecursion::Stop
1855
        )
1856
    );
1857
    rewrite_test!(
1858
        test_rewrite_f_up_stop_on_e,
1859
        transform_yes("f_down"),
1860
        transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Stop),
1861
        Transformed::new(
1862
            f_up_stop_on_e_transformed_tree(),
1863
            true,
1864
            TreeNodeRecursion::Stop
1865
        )
1866
    );
1867
1868
    transform_test!(
1869
        test_transform,
1870
        transform_yes("f_down"),
1871
        transform_yes("f_up"),
1872
        Transformed::yes(transformed_tree())
1873
    );
1874
    transform_test!(
1875
        test_transform_f_down_jump_on_a,
1876
        transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump),
1877
        transform_yes("f_up"),
1878
        Transformed::yes(transformed_tree())
1879
    );
1880
    transform_test!(
1881
        test_transform_f_down_jump_on_e,
1882
        transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump),
1883
        transform_yes("f_up"),
1884
        Transformed::yes(f_down_jump_on_e_transformed_tree())
1885
    );
1886
    transform_test!(
1887
        test_transform_f_up_jump_on_a,
1888
        transform_yes("f_down"),
1889
        transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Jump),
1890
        Transformed::yes(f_up_jump_on_a_transformed_tree())
1891
    );
1892
    transform_test!(
1893
        test_transform_f_up_jump_on_e,
1894
        transform_yes("f_down"),
1895
        transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Jump),
1896
        Transformed::yes(f_up_jump_on_e_transformed_tree())
1897
    );
1898
    transform_test!(
1899
        test_transform_f_down_stop_on_a,
1900
        transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop),
1901
        transform_yes("f_up"),
1902
        Transformed::new(
1903
            f_down_stop_on_a_transformed_tree(),
1904
            true,
1905
            TreeNodeRecursion::Stop
1906
        )
1907
    );
1908
    transform_test!(
1909
        test_transform_f_down_stop_on_e,
1910
        transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop),
1911
        transform_yes("f_up"),
1912
        Transformed::new(
1913
            f_down_stop_on_e_transformed_tree(),
1914
            true,
1915
            TreeNodeRecursion::Stop
1916
        )
1917
    );
1918
    transform_test!(
1919
        test_transform_f_up_stop_on_a,
1920
        transform_yes("f_down"),
1921
        transform_and_event_on("f_up", "f_down(a)", TreeNodeRecursion::Stop),
1922
        Transformed::new(
1923
            f_up_stop_on_a_transformed_tree(),
1924
            true,
1925
            TreeNodeRecursion::Stop
1926
        )
1927
    );
1928
    transform_test!(
1929
        test_transform_f_up_stop_on_e,
1930
        transform_yes("f_down"),
1931
        transform_and_event_on("f_up", "f_down(e)", TreeNodeRecursion::Stop),
1932
        Transformed::new(
1933
            f_up_stop_on_e_transformed_tree(),
1934
            true,
1935
            TreeNodeRecursion::Stop
1936
        )
1937
    );
1938
1939
    transform_down_test!(
1940
        test_transform_down,
1941
        transform_yes("f_down"),
1942
        Transformed::yes(transformed_down_tree())
1943
    );
1944
    transform_down_test!(
1945
        test_transform_down_f_down_jump_on_a,
1946
        transform_and_event_on("f_down", "a", TreeNodeRecursion::Jump),
1947
        Transformed::yes(f_down_jump_on_a_transformed_down_tree())
1948
    );
1949
    transform_down_test!(
1950
        test_transform_down_f_down_jump_on_e,
1951
        transform_and_event_on("f_down", "e", TreeNodeRecursion::Jump),
1952
        Transformed::yes(f_down_jump_on_e_transformed_down_tree())
1953
    );
1954
    transform_down_test!(
1955
        test_transform_down_f_down_stop_on_a,
1956
        transform_and_event_on("f_down", "a", TreeNodeRecursion::Stop),
1957
        Transformed::new(
1958
            f_down_stop_on_a_transformed_down_tree(),
1959
            true,
1960
            TreeNodeRecursion::Stop
1961
        )
1962
    );
1963
    transform_down_test!(
1964
        test_transform_down_f_down_stop_on_e,
1965
        transform_and_event_on("f_down", "e", TreeNodeRecursion::Stop),
1966
        Transformed::new(
1967
            f_down_stop_on_e_transformed_down_tree(),
1968
            true,
1969
            TreeNodeRecursion::Stop
1970
        )
1971
    );
1972
1973
    transform_up_test!(
1974
        test_transform_up,
1975
        transform_yes("f_up"),
1976
        Transformed::yes(transformed_up_tree())
1977
    );
1978
    transform_up_test!(
1979
        test_transform_up_f_up_jump_on_a,
1980
        transform_and_event_on("f_up", "a", TreeNodeRecursion::Jump),
1981
        Transformed::yes(f_up_jump_on_a_transformed_up_tree())
1982
    );
1983
    transform_up_test!(
1984
        test_transform_up_f_up_jump_on_e,
1985
        transform_and_event_on("f_up", "e", TreeNodeRecursion::Jump),
1986
        Transformed::yes(f_up_jump_on_e_transformed_up_tree())
1987
    );
1988
    transform_up_test!(
1989
        test_transform_up_f_up_stop_on_a,
1990
        transform_and_event_on("f_up", "a", TreeNodeRecursion::Stop),
1991
        Transformed::new(
1992
            f_up_stop_on_a_transformed_up_tree(),
1993
            true,
1994
            TreeNodeRecursion::Stop
1995
        )
1996
    );
1997
    transform_up_test!(
1998
        test_transform_up_f_up_stop_on_e,
1999
        transform_and_event_on("f_up", "e", TreeNodeRecursion::Stop),
2000
        Transformed::new(
2001
            f_up_stop_on_e_transformed_up_tree(),
2002
            true,
2003
            TreeNodeRecursion::Stop
2004
        )
2005
    );
2006
2007
    //             F
2008
    //          /  |  \
2009
    //       /     |     \
2010
    //    E        C        A
2011
    //    |      /   \
2012
    //    C     B     D
2013
    //  /   \         |
2014
    // B     D        A
2015
    //       |
2016
    //       A
2017
    #[test]
2018
    fn test_apply_and_visit_references() -> Result<()> {
2019
        let node_a = TestTreeNode::new(vec![], "a".to_string());
2020
        let node_b = TestTreeNode::new(vec![], "b".to_string());
2021
        let node_d = TestTreeNode::new(vec![node_a], "d".to_string());
2022
        let node_c = TestTreeNode::new(vec![node_b, node_d], "c".to_string());
2023
        let node_e = TestTreeNode::new(vec![node_c], "e".to_string());
2024
        let node_a_2 = TestTreeNode::new(vec![], "a".to_string());
2025
        let node_b_2 = TestTreeNode::new(vec![], "b".to_string());
2026
        let node_d_2 = TestTreeNode::new(vec![node_a_2], "d".to_string());
2027
        let node_c_2 = TestTreeNode::new(vec![node_b_2, node_d_2], "c".to_string());
2028
        let node_a_3 = TestTreeNode::new(vec![], "a".to_string());
2029
        let tree = TestTreeNode::new(vec![node_e, node_c_2, node_a_3], "f".to_string());
2030
2031
        let node_f_ref = &tree;
2032
        let node_e_ref = &node_f_ref.children[0];
2033
        let node_c_ref = &node_e_ref.children[0];
2034
        let node_b_ref = &node_c_ref.children[0];
2035
        let node_d_ref = &node_c_ref.children[1];
2036
        let node_a_ref = &node_d_ref.children[0];
2037
2038
        let mut m: HashMap<&TestTreeNode<String>, usize> = HashMap::new();
2039
        tree.apply(|e| {
2040
            *m.entry(e).or_insert(0) += 1;
2041
            Ok(TreeNodeRecursion::Continue)
2042
        })?;
2043
2044
        let expected = HashMap::from([
2045
            (node_f_ref, 1),
2046
            (node_e_ref, 1),
2047
            (node_c_ref, 2),
2048
            (node_d_ref, 2),
2049
            (node_b_ref, 2),
2050
            (node_a_ref, 3),
2051
        ]);
2052
        assert_eq!(m, expected);
2053
2054
        struct TestVisitor<'n> {
2055
            m: HashMap<&'n TestTreeNode<String>, (usize, usize)>,
2056
        }
2057
2058
        impl<'n> TreeNodeVisitor<'n> for TestVisitor<'n> {
2059
            type Node = TestTreeNode<String>;
2060
2061
            fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
2062
                let (down_count, _) = self.m.entry(node).or_insert((0, 0));
2063
                *down_count += 1;
2064
                Ok(TreeNodeRecursion::Continue)
2065
            }
2066
2067
            fn f_up(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
2068
                let (_, up_count) = self.m.entry(node).or_insert((0, 0));
2069
                *up_count += 1;
2070
                Ok(TreeNodeRecursion::Continue)
2071
            }
2072
        }
2073
2074
        let mut visitor = TestVisitor { m: HashMap::new() };
2075
        tree.visit(&mut visitor)?;
2076
2077
        let expected = HashMap::from([
2078
            (node_f_ref, (1, 1)),
2079
            (node_e_ref, (1, 1)),
2080
            (node_c_ref, (2, 2)),
2081
            (node_d_ref, (2, 2)),
2082
            (node_b_ref, (2, 2)),
2083
            (node_a_ref, (3, 3)),
2084
        ]);
2085
        assert_eq!(visitor.m, expected);
2086
2087
        Ok(())
2088
    }
2089
}