vortex_array/expr/traversal/
visitor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::marker::PhantomData;
5
6use vortex_error::VortexResult;
7
8use crate::expr::traversal::Node;
9use crate::expr::traversal::NodeExt;
10use crate::expr::traversal::NodeVisitor;
11use crate::expr::traversal::TraversalOrder;
12
13struct FnVisitor<'a, F, T: 'a>
14where
15    F: FnMut(&'a T) -> VortexResult<TraversalOrder>,
16{
17    f_down: Option<F>,
18    f_up: Option<F>,
19    _data: PhantomData<&'a T>,
20}
21
22impl<'a, T, F> NodeVisitor<'a> for FnVisitor<'a, F, T>
23where
24    F: FnMut(&'a T) -> VortexResult<TraversalOrder>,
25    T: NodeExt,
26{
27    type NodeTy = T;
28
29    fn visit_down(&mut self, node: &'a T) -> VortexResult<TraversalOrder> {
30        if let Some(f) = self.f_down.as_mut() {
31            f(node)
32        } else {
33            Ok(TraversalOrder::Continue)
34        }
35    }
36
37    fn visit_up(&mut self, node: &'a T) -> VortexResult<TraversalOrder> {
38        if let Some(f) = self.f_up.as_mut() {
39            f(node)
40        } else {
41            Ok(TraversalOrder::Continue)
42        }
43    }
44}
45
46/// Traverse a [`Node`]-based tree using a closure. It will do it by walking the tree from the bottom going up.
47pub fn pre_order_visit_up<'a, T: 'a + Node>(
48    tree: &'a T,
49    f: impl FnMut(&'a T) -> VortexResult<TraversalOrder>,
50) -> VortexResult<()> {
51    let mut visitor = FnVisitor {
52        f_down: None,
53        f_up: Some(f),
54        _data: PhantomData,
55    };
56
57    tree.accept(&mut visitor)?;
58
59    Ok(())
60}
61
62/// Traverse a [`Node`]-based tree using a closure. It will do it by walking the tree from the top going down.
63pub fn pre_order_visit_down<'a, T: 'a + Node>(
64    tree: &'a T,
65    f: impl FnMut(&'a T) -> VortexResult<TraversalOrder>,
66) -> VortexResult<()> {
67    let mut visitor = FnVisitor {
68        f_down: Some(f),
69        f_up: None,
70        _data: PhantomData,
71    };
72
73    tree.accept(&mut visitor)?;
74
75    Ok(())
76}
vortex_array/expr/traversal/visitor.rs

vortex_array/expr/traversal/
visitor.rs
