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

vortex_array/expr/traversal/
visitor.rs
