Skip to main content

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: F,
18    _data: PhantomData<&'a T>,
19}
20
21impl<'a, T, F> NodeVisitor<'a> for FnVisitor<'a, F, T>
22where
23    F: FnMut(&'a T) -> VortexResult<TraversalOrder>,
24    T: NodeExt,
25{
26    type NodeTy = T;
27
28    fn visit_down(&mut self, node: &'a T) -> VortexResult<TraversalOrder> {
29        (self.f)(node)
30    }
31}
32
33/// Traverse a [`Node`]-based tree using a closure. It will do it by walking the tree from the top going down.
34pub fn pre_order_visit_down<'a, T: 'a + Node>(
35    tree: &'a T,
36    f: impl FnMut(&'a T) -> VortexResult<TraversalOrder>,
37) -> VortexResult<()> {
38    let mut visitor = FnVisitor {
39        f,
40        _data: PhantomData,
41    };
42
43    tree.accept(&mut visitor)?;
44
45    Ok(())
46}