/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 | | } |