16440250234

Committed 22 Jul 2025 09:25AM UTC coverage: 81.473% (-0.07%) from 81.54%

Build # 16440250234

Build Type

push

github

Committed by

web-flow

Commit Message

feat: Refactor traversal logic (#3926)

The main rationale for this change is that I would like to use this
logic elsewhere, but it wasn't touched in a while and needed some
polishing.

I've moved most of the code that deals with actually applying operations
in a certain order into the default implementation of `Node::accept`,
`Node::rewrite` and the `transform_` functions.
A future user of `Node` will only need to fill in the structure of their
tree, getting the fiddly bits OOTB.

This is not a fully-fledged solution for any possible usecase, but its
an attempt at tightening the existing abstractions and leaving something
that could be extended to new or more advanced patterns.

---------

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Run Details

254 of 298 new or added lines in 8 files covered. (85.23%)

8 existing lines in 3 files now uncovered.

41953 of 51493 relevant lines covered (81.47%)

171685.53 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

87.96

/vortex-expr/src/traversal/mod.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Datafusion inspired tree traversal logic.
//!
//! Users should want to implement [`Node`] and potentially [`NodeContainer`].

mod references;
mod visitor;

use std::marker::PhantomData;

use itertools::Itertools;
pub use references::ReferenceCollector;
pub use visitor::{pre_order_visit_down, pre_order_visit_up};
use vortex_error::VortexResult;

use crate::ExprRef;

/// Signal to control a traversal's flow
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TraversalOrder {
    /// In a top-down traversal, skip visiting the children of the current node.
    /// In the bottom-up phase of the traversal, skip the next step. Either skipping the children of the node,
    /// moving to its next sibling, or skipping its parent once the children are traversed.
    Skip,
    /// Stop visiting any more nodes in the traversal.
    Stop,
    /// Continue with the traversal as expected.
    Continue,
}

impl TraversalOrder {
    /// If directed to, continue to visit nodes by running `f`, which should apply on the node's children.
    pub fn visit_children<F: FnOnce() -> VortexResult<TraversalOrder>>(
        self,
        f: F,
    ) -> VortexResult<TraversalOrder> {
        match self {
            Self::Skip => Ok(TraversalOrder::Continue),
            Self::Stop => Ok(self),
            Self::Continue => f(),
        }
    }

    /// If directed to, continue to visit nodes by running `f`, which should apply on the node's parent.
    pub fn visit_parent<F: FnOnce() -> VortexResult<TraversalOrder>>(
        self,
        f: F,
    ) -> VortexResult<TraversalOrder> {
        match self {
            Self::Continue => f(),
            Self::Skip | Self::Stop => Ok(self),
        }
    }
}

#[derive(Debug, Clone)]
pub struct Transformed<T> {
    pub value: T,
    pub order: TraversalOrder,
    pub changed: bool,
}

impl<T> Transformed<T> {
    pub fn yes(value: T) -> Self {
        Self {
            value,
            order: TraversalOrder::Continue,
            changed: true,
        }
    }

    pub fn no(value: T) -> Self {
        Self {
            value,
            order: TraversalOrder::Continue,
            changed: false,
        }
    }

    pub fn into_inner(self) -> T {
        self.value
    }
}

pub trait NodeVisitor<'a> {
    type NodeTy: Node;

    fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
        _ = node;
        Ok(TraversalOrder::Continue)
    }

    fn visit_up(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {
        _ = node;
        Ok(TraversalOrder::Continue)
    }
}

pub trait NodeRewriter: Sized {
    type NodeTy: Node;

    fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
        Ok(Transformed::no(node))
    }

    fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
        Ok(Transformed::no(node))
    }
}

pub trait Node: Sized + Clone {
    /// Walk the node's children by applying `f` to them.
    ///
    /// This is a lower level API that other functions rely on for correctness.
    fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
        &'a self,
        f: F,
    ) -> VortexResult<TraversalOrder>;

    /// Rewrite the node's children by applying `f` to them.
    ///
    /// This is a lower level API that other functions rely on for correctness.
    fn map_children<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
        self,
        f: F,
    ) -> VortexResult<Transformed<Self>>;

    /// Walk the tree in pre-order (top-down) way, rewriting it as it goes.
    fn rewrite<R: NodeRewriter<NodeTy = Self>>(
        self,
        rewriter: &mut R,
    ) -> VortexResult<Transformed<Self>> {
        let mut transformed = rewriter.visit_down(self)?;

        let transformed = match transformed.order {
            TraversalOrder::Stop => Ok(transformed),
            TraversalOrder::Skip => {
                transformed.order = TraversalOrder::Continue;
                Ok(transformed)
            }
            TraversalOrder::Continue => transformed
                .value
                .map_children(|c| c.rewrite(rewriter))
                .map(|mut t| {
                    t.changed |= transformed.changed;
                    t
                }),
        }?;

        match transformed.order {
            TraversalOrder::Stop | TraversalOrder::Skip => Ok(transformed),
            TraversalOrder::Continue => {
                let mut up_rewrite = rewriter.visit_up(transformed.value)?;
                up_rewrite.changed |= transformed.changed;
                Ok(up_rewrite)
            }
        }
    }

    /// A pre-order (top-down) traversal.
    fn accept<'a, V: NodeVisitor<'a, NodeTy = Self>>(
        &'a self,
        visitor: &mut V,
    ) -> VortexResult<TraversalOrder> {
        visitor
            .visit_down(self)?
            .visit_children(|| self.apply_children(|c| c.accept(visitor)))?
            .visit_parent(|| visitor.visit_up(self))
    }

    /// A pre-order transformation
    fn transform_down<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
        self,
        f: F,
    ) -> VortexResult<Transformed<Self>> {
        let mut rewriter = FnRewriter {
            f_down: Some(f),
            f_up: None,
            _data: PhantomData,
        };

        self.rewrite(&mut rewriter)
    }

    /// A post-order transform
    fn transform_up<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
        self,
        f: F,
    ) -> VortexResult<Transformed<Self>> {
        let mut rewriter = FnRewriter {
            f_down: None,
            f_up: Some(f),
            _data: PhantomData,
        };

        self.rewrite(&mut rewriter)
    }
}

struct FnRewriter<F, T> {
    f_down: Option<F>,
    f_up: Option<F>,
    _data: PhantomData<T>,
}

impl<F, T> NodeRewriter for FnRewriter<F, T>
where
    T: Node,
    F: FnMut(T) -> VortexResult<Transformed<T>>,
{
    type NodeTy = T;

    fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
        if let Some(f) = self.f_down.as_mut() {
            f(node)
        } else {
            Ok(Transformed::no(node))
        }
    }

    fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
        if let Some(f) = self.f_up.as_mut() {
            f(node)
        } else {
            Ok(Transformed::no(node))
        }
    }
}

/// A container holding a [`Node`]'s children, which a function can be applied (or mapped) to.
///
/// The trait is also implemented to container types in order to make implementing [`Node::map_children`]
/// and [`Node::apply_children`] easier.
pub trait NodeContainer<'a, T: 'a>: Sized {
    /// Applies `f` to all elements of the container, accepting them by reference
    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
        &'a self,
        f: F,
    ) -> VortexResult<TraversalOrder>;

    /// Consumes all the children of the node, replacing them with the result of `f`.
    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
        self,
        f: F,
    ) -> VortexResult<Transformed<Self>>;
}

pub trait NodeRefContainer<'a, T: 'a>: Sized {
    fn apply_ref_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
        &self,
        f: F,
    ) -> VortexResult<TraversalOrder>;
}

impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeRefContainer<'a, T> for Vec<&'a C> {
    fn apply_ref_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
        &self,
        mut f: F,
    ) -> VortexResult<TraversalOrder> {
        let mut order = TraversalOrder::Continue;

        for c in self {
            order = c.apply_elements(&mut f)?;
            match order {
                TraversalOrder::Continue | TraversalOrder::Skip => {}
                TraversalOrder::Stop => return Ok(TraversalOrder::Stop),
            }
        }

        Ok(order)
    }
}

impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeContainer<'a, T> for Vec<C> {
    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
        &'a self,
        mut f: F,
    ) -> VortexResult<TraversalOrder> {
        let mut order = TraversalOrder::Continue;

        for c in self {
            order = c.apply_elements(&mut f)?;
            match order {
                TraversalOrder::Continue | TraversalOrder::Skip => {}
                TraversalOrder::Stop => return Ok(TraversalOrder::Stop),
            }
        }

        Ok(order)
    }

    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
        self,
        mut f: F,
    ) -> VortexResult<Transformed<Self>> {
        let mut order = TraversalOrder::Continue;
        let mut changed = false;

        let value = self
            .into_iter()
            .map(|c| match order {
                TraversalOrder::Continue | TraversalOrder::Skip => {
                    c.map_elements(&mut f).map(|result| {
                        order = result.order;
                        changed |= result.changed;
                        result.value
                    })
                }
                TraversalOrder::Stop => Ok(c),
            })
            .collect::<VortexResult<Vec<_>>>()?;

        Ok(Transformed {
            value,
            order,
            changed,
        })
    }
}

impl<'a> NodeContainer<'a, Self> for ExprRef {
    fn apply_elements<F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
        &'a self,
        mut f: F,
    ) -> VortexResult<TraversalOrder> {
        f(self)
    }

    fn map_elements<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
        self,
        mut f: F,
    ) -> VortexResult<Transformed<Self>> {
        f(self)
    }
}

impl Node for ExprRef {
    fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
        &'a self,
        mut f: F,
    ) -> VortexResult<TraversalOrder> {
        self.children().apply_ref_elements(&mut f)
    }

    fn map_children<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
        self,
        f: F,
    ) -> VortexResult<Transformed<Self>> {
        let transformed = self
            .children()
            .into_iter()
            .cloned()
            .collect_vec()
            .map_elements(f)?;

        if transformed.changed {
            Ok(Transformed {
                value: self.with_children(transformed.value)?,
                order: transformed.order,
                changed: true,
            })
        } else {
            Ok(Transformed::no(self))
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use vortex_error::VortexResult;
    use vortex_utils::aliases::hash_set::HashSet;

    use crate::traversal::visitor::pre_order_visit_down;
    use crate::traversal::{Node, NodeRewriter, NodeVisitor, Transformed, TraversalOrder};
    use crate::{
        BinaryExpr, BinaryVTable, ExprRef, GetItemVTable, IntoExpr, LiteralExpr, LiteralVTable,
        Operator, VortexExpr, col, is_root, root,
    };

    #[derive(Default)]
    pub struct ExprLitCollector<'a>(pub Vec<&'a ExprRef>);

    impl<'a> NodeVisitor<'a> for ExprLitCollector<'a> {
        type NodeTy = ExprRef;

        fn visit_down(&mut self, node: &'a ExprRef) -> VortexResult<TraversalOrder> {
            if node.is::<LiteralVTable>() {
                self.0.push(node)
            }
            Ok(TraversalOrder::Continue)
        }

        fn visit_up(&mut self, _node: &'a ExprRef) -> VortexResult<TraversalOrder> {
            Ok(TraversalOrder::Continue)
        }
    }

    fn expr_col_to_lit_transform(
        node: ExprRef,
        idx: &mut i32,
    ) -> VortexResult<Transformed<ExprRef>> {
        if node.is::<GetItemVTable>() {
            let lit_id = *idx;
            *idx += 1;
            Ok(Transformed::yes(LiteralExpr::new_expr(lit_id)))
        } else {
            Ok(Transformed::no(node))
        }
    }

    #[derive(Default)]
    pub struct SkipDownRewriter;

    impl NodeRewriter for SkipDownRewriter {
        type NodeTy = ExprRef;

        fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
            Ok(Transformed {
                value: node,
                order: TraversalOrder::Skip,
                changed: false,
            })
        }

        fn visit_up(&mut self, _node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {
            Ok(Transformed::yes(root()))
        }
    }

    #[test]
    fn expr_deep_visitor_test() {
        let col1: Arc<dyn VortexExpr> = col("col1");
        let lit1 = LiteralExpr::new(1).into_expr();
        let expr = BinaryExpr::new(col1.clone(), Operator::Eq, lit1.clone()).into_expr();
        let lit2 = LiteralExpr::new(2).into_expr();
        let expr = BinaryExpr::new(expr, Operator::And, lit2).into_expr();
        let mut printer = ExprLitCollector::default();
        expr.accept(&mut printer).unwrap();
        assert_eq!(printer.0.len(), 2);
    }

    #[test]
    fn expr_deep_mut_visitor_test() {
        let col1: Arc<dyn VortexExpr> = col("col1");
        let col2: Arc<dyn VortexExpr> = col("col2");
        let expr = BinaryExpr::new_expr(col1.clone(), Operator::Eq, col2.clone());
        let lit2 = LiteralExpr::new_expr(2);
        let expr = BinaryExpr::new_expr(expr, Operator::And, lit2);

        let mut idx = 0_i32;
        let new = expr
            .transform_up(|node| expr_col_to_lit_transform(node, &mut idx))
            .unwrap();
        assert!(new.changed);

        let expr = new.value;

        let mut printer = ExprLitCollector::default();
        expr.accept(&mut printer).unwrap();
        assert_eq!(printer.0.len(), 3);
    }

    #[test]
    fn expr_skip_test() {
        let col1: Arc<dyn VortexExpr> = col("col1");
        let col2: Arc<dyn VortexExpr> = col("col2");
        let expr1 = BinaryExpr::new_expr(col1.clone(), Operator::Eq, col2.clone());
        let col3: Arc<dyn VortexExpr> = col("col3");
        let col4: Arc<dyn VortexExpr> = col("col4");
        let expr2 = BinaryExpr::new_expr(col3.clone(), Operator::NotEq, col4.clone());
        let expr = BinaryExpr::new_expr(expr1, Operator::And, expr2);

        let mut nodes = Vec::new();
        pre_order_visit_down(&expr, |node: &ExprRef| {
            if node.is::<GetItemVTable>() {
                nodes.push(node)
            }
            if let Some(bin) = node.as_opt::<BinaryVTable>() {
                if bin.op() == Operator::Eq {
                    return Ok(TraversalOrder::Skip);
                }
            }
            Ok(TraversalOrder::Continue)
        })
        .unwrap();

        let nodes: HashSet<ExprRef> = HashSet::from_iter(nodes.into_iter().cloned());
        assert_eq!(nodes, HashSet::from_iter([col("col3"), col("col4")]));
    }

    #[test]
    fn expr_stop_test() {
        let col1: Arc<dyn VortexExpr> = col("col1");
        let col2: Arc<dyn VortexExpr> = col("col2");
        let expr1 = BinaryExpr::new_expr(col1.clone(), Operator::Eq, col2.clone());
        let col3: Arc<dyn VortexExpr> = col("col3");
        let col4: Arc<dyn VortexExpr> = col("col4");
        let expr2 = BinaryExpr::new_expr(col3.clone(), Operator::NotEq, col4.clone());
        let expr = BinaryExpr::new_expr(expr1, Operator::And, expr2);

        let mut nodes = Vec::new();
        pre_order_visit_down(&expr, |node: &ExprRef| {
            if node.is::<GetItemVTable>() {
                nodes.push(node)
            }
            if let Some(bin) = node.as_opt::<BinaryVTable>() {
                if bin.op() == Operator::Eq {
                    return Ok(TraversalOrder::Stop);
                }
            }
            Ok(TraversalOrder::Continue)
        })
        .unwrap();

        assert!(nodes.is_empty());
    }

    #[test]
    fn expr_skip_down_visit_up() {
        let col = col("col");

        let mut visitor = SkipDownRewriter;
        let result = col.rewrite(&mut visitor).unwrap();

        assert!(result.changed);
        assert!(is_root(&result.value));
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	//! Datafusion inspired tree traversal logic.
5	//!
6	//! Users should want to implement [`Node`] and potentially [`NodeContainer`].
7
8	mod references;
9	mod visitor;
10
11	use std::marker::PhantomData;
12
13	use itertools::Itertools;
14	pub use references::ReferenceCollector;
15	pub use visitor::{pre_order_visit_down, pre_order_visit_up};
16	use vortex_error::VortexResult;
17
18	use crate::ExprRef;
19
20	/// Signal to control a traversal's flow
21	#[derive(Debug, Clone, PartialEq, Eq)]
22	pub enum TraversalOrder {
23	/// In a top-down traversal, skip visiting the children of the current node.
24	/// In the bottom-up phase of the traversal, skip the next step. Either skipping the children of the node,
25	/// moving to its next sibling, or skipping its parent once the children are traversed.
26	Skip,
27	/// Stop visiting any more nodes in the traversal.
28	Stop,
29	/// Continue with the traversal as expected.
30	Continue,
31	}
32
33	impl TraversalOrder {
34	/// If directed to, continue to visit nodes by running `f`, which should apply on the node's children.
35	pub fn visit_children<F: FnOnce() -> VortexResult<TraversalOrder>>(	62,668✔
36	self,	62,668✔
37	f: F,	62,668✔
38	) -> VortexResult<TraversalOrder> {	62,668✔
39	match self {	62,668✔
40	Self::Skip => Ok(TraversalOrder::Continue),	47,233✔
41	Self::Stop => Ok(self),	3✔
42	Self::Continue => f(),	15,432✔
43	}
44	}	62,668✔
45
46	/// If directed to, continue to visit nodes by running `f`, which should apply on the node's parent.
47	pub fn visit_parent<F: FnOnce() -> VortexResult<TraversalOrder>>(	62,668✔
48	self,	62,668✔
49	f: F,	62,668✔
50	) -> VortexResult<TraversalOrder> {	62,668✔
51	match self {	62,668✔
52	Self::Continue => f(),	62,664✔
53	Self::Skip \| Self::Stop => Ok(self),	4✔
54	}
55	}	62,668✔
56	}
57
58	#[derive(Debug, Clone)]
59	pub struct Transformed<T> {
60	pub value: T,
61	pub order: TraversalOrder,
62	pub changed: bool,
63	}
64
65	impl<T> Transformed<T> {
66	pub fn yes(value: T) -> Self {	13,330✔
67	Self {	13,330✔
68	value,	13,330✔
69	order: TraversalOrder::Continue,	13,330✔
70	changed: true,	13,330✔
71	}	13,330✔
72	}	13,330✔
73
74	pub fn no(value: T) -> Self {	1,539,824✔
75	Self {	1,539,824✔
76	value,	1,539,824✔
77	order: TraversalOrder::Continue,	1,539,824✔
78	changed: false,	1,539,824✔
79	}	1,539,824✔
80	}	1,539,824✔
81
82	pub fn into_inner(self) -> T {	71,388✔
83	self.value	71,388✔
84	}	71,388✔
85	}
86
87	pub trait NodeVisitor<'a> {
88	type NodeTy: Node;
89
NEW 90	fn visit_down(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {	×
NEW 91	_ = node;	×
92	Ok(TraversalOrder::Continue)	×
93	}	×
94
95	fn visit_up(&mut self, node: &'a Self::NodeTy) -> VortexResult<TraversalOrder> {	36,803✔
96	_ = node;	36,803✔
97	Ok(TraversalOrder::Continue)	36,803✔
98	}	36,803✔
99	}
100
101	pub trait NodeRewriter: Sized {
102	type NodeTy: Node;
103
NEW 104	fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {	×
NEW 105	Ok(Transformed::no(node))	×
UNCOV 106	}	×
107
108	fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {	147,534✔
109	Ok(Transformed::no(node))	147,534✔
110	}	147,534✔
111	}
112
113	pub trait Node: Sized + Clone {
114	/// Walk the node's children by applying `f` to them.
115	///
116	/// This is a lower level API that other functions rely on for correctness.
117	fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(
118	&'a self,
119	f: F,
120	) -> VortexResult<TraversalOrder>;
121
122	/// Rewrite the node's children by applying `f` to them.
123	///
124	/// This is a lower level API that other functions rely on for correctness.
125	fn map_children<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(
126	self,
127	f: F,
128	) -> VortexResult<Transformed<Self>>;
129
130	/// Walk the tree in pre-order (top-down) way, rewriting it as it goes.
131	fn rewrite<R: NodeRewriter<NodeTy = Self>>(	525,767✔
132	self,	525,767✔
133	rewriter: &mut R,	525,767✔
134	) -> VortexResult<Transformed<Self>> {	525,767✔
135	let mut transformed = rewriter.visit_down(self)?;	525,767✔
136
137	let transformed = match transformed.order {	525,767✔
NEW 138	TraversalOrder::Stop => Ok(transformed),	×
139	TraversalOrder::Skip => {
140	transformed.order = TraversalOrder::Continue;	6,088✔
141	Ok(transformed)	6,088✔
142	}
143	TraversalOrder::Continue => transformed	519,679✔
144	.value	519,679✔
145	.map_children(\|c\| c.rewrite(rewriter))	519,679✔
146	.map(\|mut t\| {	519,679✔
147	t.changed \|= transformed.changed;	519,679✔
148	t	519,679✔
149	}),	519,679✔
NEW 150	}?;	×
151
152	match transformed.order {	525,767✔
NEW 153	TraversalOrder::Stop \| TraversalOrder::Skip => Ok(transformed),	×
154	TraversalOrder::Continue => {
155	let mut up_rewrite = rewriter.visit_up(transformed.value)?;	525,767✔
156	up_rewrite.changed \|= transformed.changed;	525,767✔
157	Ok(up_rewrite)	525,767✔
158	}
159	}
160	}	525,767✔
161
162	/// A pre-order (top-down) traversal.
163	fn accept<'a, V: NodeVisitor<'a, NodeTy = Self>>(	62,668✔
164	&'a self,	62,668✔
165	visitor: &mut V,	62,668✔
166	) -> VortexResult<TraversalOrder> {	62,668✔
167	visitor	62,668✔
168	.visit_down(self)?	62,668✔
169	.visit_children(\|\| self.apply_children(\|c\| c.accept(visitor)))?	62,668✔
170	.visit_parent(\|\| visitor.visit_up(self))	62,668✔
171	}	62,668✔
172
173	/// A pre-order transformation
NEW 174	fn transform_down<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(	×
NEW 175	self,	×
NEW 176	f: F,	×
NEW 177	) -> VortexResult<Transformed<Self>> {	×
NEW 178	let mut rewriter = FnRewriter {	×
NEW 179	f_down: Some(f),	×
NEW 180	f_up: None,	×
NEW 181	_data: PhantomData,	×
NEW 182	};	×
183
NEW 184	self.rewrite(&mut rewriter)	×
UNCOV 185	}	×
186
187	/// A post-order transform
188	fn transform_up<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(	71,389✔
189	self,	71,389✔
190	f: F,	71,389✔
191	) -> VortexResult<Transformed<Self>> {	71,389✔
192	let mut rewriter = FnRewriter {	71,389✔
193	f_down: None,	71,389✔
194	f_up: Some(f),	71,389✔
195	_data: PhantomData,	71,389✔
196	};	71,389✔
197
198	self.rewrite(&mut rewriter)	71,389✔
199	}	71,389✔
200	}
201
202	struct FnRewriter<F, T> {
203	f_down: Option<F>,
204	f_up: Option<F>,
205	_data: PhantomData<T>,
206	}
207
208	impl<F, T> NodeRewriter for FnRewriter<F, T>
209	where
210	T: Node,
211	F: FnMut(T) -> VortexResult<Transformed<T>>,
212	{
213	type NodeTy = T;
214
215	fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {	371,197✔
216	if let Some(f) = self.f_down.as_mut() {	371,197✔
NEW 217	f(node)	×
218	} else {
219	Ok(Transformed::no(node))	371,197✔
220	}
221	}	371,197✔
222
223	fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {	371,197✔
224	if let Some(f) = self.f_up.as_mut() {	371,197✔
225	f(node)	371,197✔
226	} else {
NEW 227	Ok(Transformed::no(node))	×
228	}
229	}	371,197✔
230	}
231
232	/// A container holding a [`Node`]'s children, which a function can be applied (or mapped) to.
233	///
234	/// The trait is also implemented to container types in order to make implementing [`Node::map_children`]
235	/// and [`Node::apply_children`] easier.
236	pub trait NodeContainer<'a, T: 'a>: Sized {
237	/// Applies `f` to all elements of the container, accepting them by reference
238	fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
239	&'a self,
240	f: F,
241	) -> VortexResult<TraversalOrder>;
242
243	/// Consumes all the children of the node, replacing them with the result of `f`.
244	fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
245	self,
246	f: F,
247	) -> VortexResult<Transformed<Self>>;
248	}
249
250	pub trait NodeRefContainer<'a, T: 'a>: Sized {
251	fn apply_ref_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
252	&self,
253	f: F,
254	) -> VortexResult<TraversalOrder>;
255	}
256
257	impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeRefContainer<'a, T> for Vec<&'a C> {
258	fn apply_ref_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(	15,432✔
259	&self,	15,432✔
260	mut f: F,	15,432✔
261	) -> VortexResult<TraversalOrder> {	15,432✔
262	let mut order = TraversalOrder::Continue;	15,432✔
263
264	for c in self {	36,729✔
265	order = c.apply_elements(&mut f)?;	21,298✔
266	match order {	21,298✔
267	TraversalOrder::Continue \| TraversalOrder::Skip => {}	21,297✔
268	TraversalOrder::Stop => return Ok(TraversalOrder::Stop),	1✔
269	}
270	}
271
272	Ok(order)	15,431✔
273	}	15,432✔
274	}
275
276	impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeContainer<'a, T> for Vec<C> {
NEW 277	fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(	×
UNCOV 278	&'a self,	×
NEW 279	mut f: F,	×
NEW 280	) -> VortexResult<TraversalOrder> {	×
NEW 281	let mut order = TraversalOrder::Continue;	×
282
NEW 283	for c in self {	×
NEW 284	order = c.apply_elements(&mut f)?;	×
NEW 285	match order {	×
NEW 286	TraversalOrder::Continue \| TraversalOrder::Skip => {}	×
NEW 287	TraversalOrder::Stop => return Ok(TraversalOrder::Stop),	×
288	}
289	}
290
NEW 291	Ok(order)	×
UNCOV 292	}	×
293
294	fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(	519,679✔
295	self,	519,679✔
296	mut f: F,	519,679✔
297	) -> VortexResult<Transformed<Self>> {	519,679✔
298	let mut order = TraversalOrder::Continue;	519,679✔
299	let mut changed = false;	519,679✔
300
301	let value = self	519,679✔
302	.into_iter()	519,679✔
303	.map(\|c\| match order {	519,679✔
304	TraversalOrder::Continue \| TraversalOrder::Skip => {
305	c.map_elements(&mut f).map(\|result\| {	417,377✔
306	order = result.order;	417,377✔
307	changed \|= result.changed;	417,377✔
308	result.value	417,377✔
309	})	417,377✔
310	}
NEW 311	TraversalOrder::Stop => Ok(c),	×
312	})	417,377✔
313	.collect::<VortexResult<Vec<_>>>()?;	519,679✔
314
315	Ok(Transformed {	519,679✔
316	value,	519,679✔
317	order,	519,679✔
318	changed,	519,679✔
319	})	519,679✔
320	}	519,679✔
321	}
322
323	impl<'a> NodeContainer<'a, Self> for ExprRef {
324	fn apply_elements<F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(	21,298✔
325	&'a self,	21,298✔
326	mut f: F,	21,298✔
327	) -> VortexResult<TraversalOrder> {	21,298✔
328	f(self)	21,298✔
329	}	21,298✔
330
331	fn map_elements<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(	417,377✔
332	self,	417,377✔
333	mut f: F,	417,377✔
334	) -> VortexResult<Transformed<Self>> {	417,377✔
335	f(self)	417,377✔
336	}	417,377✔
337	}
338
339	impl Node for ExprRef {
340	fn apply_children<'a, F: FnMut(&'a Self) -> VortexResult<TraversalOrder>>(	15,432✔
341	&'a self,	15,432✔
342	mut f: F,	15,432✔
343	) -> VortexResult<TraversalOrder> {	15,432✔
344	self.children().apply_ref_elements(&mut f)	15,432✔
345	}	15,432✔
346
347	fn map_children<F: FnMut(Self) -> VortexResult<Transformed<Self>>>(	519,679✔
348	self,	519,679✔
349	f: F,	519,679✔
350	) -> VortexResult<Transformed<Self>> {	519,679✔
351	let transformed = self	519,679✔
352	.children()	519,679✔
353	.into_iter()	519,679✔
354	.cloned()	519,679✔
355	.collect_vec()	519,679✔
356	.map_elements(f)?;	519,679✔
357
358	if transformed.changed {	519,679✔
359	Ok(Transformed {
360	value: self.with_children(transformed.value)?,	12,845✔
361	order: transformed.order,	12,845✔
362	changed: true,
363	})
364	} else {
365	Ok(Transformed::no(self))	506,834✔
366	}
367	}	519,679✔
368	}
369
370	#[cfg(test)]
371	mod tests {
372	use std::sync::Arc;
373
374	use vortex_error::VortexResult;
375	use vortex_utils::aliases::hash_set::HashSet;
376
377	use crate::traversal::visitor::pre_order_visit_down;
378	use crate::traversal::{Node, NodeRewriter, NodeVisitor, Transformed, TraversalOrder};
379	use crate::{
380	BinaryExpr, BinaryVTable, ExprRef, GetItemVTable, IntoExpr, LiteralExpr, LiteralVTable,
381	Operator, VortexExpr, col, is_root, root,
382	};
383
384	#[derive(Default)]
385	pub struct ExprLitCollector<'a>(pub Vec<&'a ExprRef>);
386
387	impl<'a> NodeVisitor<'a> for ExprLitCollector<'a> {
388	type NodeTy = ExprRef;
389
390	fn visit_down(&mut self, node: &'a ExprRef) -> VortexResult<TraversalOrder> {	11✔
391	if node.is::<LiteralVTable>() {	11✔
392	self.0.push(node)	5✔
393	}	6✔
394	Ok(TraversalOrder::Continue)	11✔
395	}	11✔
396
397	fn visit_up(&mut self, _node: &'a ExprRef) -> VortexResult<TraversalOrder> {	11✔
398	Ok(TraversalOrder::Continue)	11✔
399	}	11✔
400	}
401
402	fn expr_col_to_lit_transform(	7✔
403	node: ExprRef,	7✔
404	idx: &mut i32,	7✔
405	) -> VortexResult<Transformed<ExprRef>> {	7✔
406	if node.is::<GetItemVTable>() {	7✔
407	let lit_id = *idx;	2✔
408	*idx += 1;	2✔
409	Ok(Transformed::yes(LiteralExpr::new_expr(lit_id)))	2✔
410	} else {
411	Ok(Transformed::no(node))	5✔
412	}
413	}	7✔
414
415	#[derive(Default)]
416	pub struct SkipDownRewriter;
417
418	impl NodeRewriter for SkipDownRewriter {
419	type NodeTy = ExprRef;
420
421	fn visit_down(&mut self, node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {	1✔
422	Ok(Transformed {	1✔
423	value: node,	1✔
424	order: TraversalOrder::Skip,	1✔
425	changed: false,	1✔
426	})	1✔
427	}	1✔
428
429	fn visit_up(&mut self, _node: Self::NodeTy) -> VortexResult<Transformed<Self::NodeTy>> {	1✔
430	Ok(Transformed::yes(root()))	1✔
431	}	1✔
432	}
433
434	#[test]
435	fn expr_deep_visitor_test() {	1✔
436	let col1: Arc<dyn VortexExpr> = col("col1");	1✔
437	let lit1 = LiteralExpr::new(1).into_expr();	1✔
438	let expr = BinaryExpr::new(col1.clone(), Operator::Eq, lit1.clone()).into_expr();	1✔
439	let lit2 = LiteralExpr::new(2).into_expr();	1✔
440	let expr = BinaryExpr::new(expr, Operator::And, lit2).into_expr();	1✔
441	let mut printer = ExprLitCollector::default();	1✔
442	expr.accept(&mut printer).unwrap();	1✔
443	assert_eq!(printer.0.len(), 2);	1✔
444	}	1✔
445
446	#[test]
447	fn expr_deep_mut_visitor_test() {	1✔
448	let col1: Arc<dyn VortexExpr> = col("col1");	1✔
449	let col2: Arc<dyn VortexExpr> = col("col2");	1✔
450	let expr = BinaryExpr::new_expr(col1.clone(), Operator::Eq, col2.clone());	1✔
451	let lit2 = LiteralExpr::new_expr(2);	1✔
452	let expr = BinaryExpr::new_expr(expr, Operator::And, lit2);	1✔
453
454	let mut idx = 0_i32;	1✔
455	let new = expr	1✔
456	.transform_up(\|node\| expr_col_to_lit_transform(node, &mut idx))	7✔
457	.unwrap();	1✔
458	assert!(new.changed);	1✔
459
460	let expr = new.value;	1✔
461
462	let mut printer = ExprLitCollector::default();	1✔
463	expr.accept(&mut printer).unwrap();	1✔
464	assert_eq!(printer.0.len(), 3);	1✔
465	}	1✔
466
467	#[test]
468	fn expr_skip_test() {	1✔
469	let col1: Arc<dyn VortexExpr> = col("col1");	1✔
470	let col2: Arc<dyn VortexExpr> = col("col2");	1✔
471	let expr1 = BinaryExpr::new_expr(col1.clone(), Operator::Eq, col2.clone());	1✔
472	let col3: Arc<dyn VortexExpr> = col("col3");	1✔
473	let col4: Arc<dyn VortexExpr> = col("col4");	1✔
474	let expr2 = BinaryExpr::new_expr(col3.clone(), Operator::NotEq, col4.clone());	1✔
475	let expr = BinaryExpr::new_expr(expr1, Operator::And, expr2);	1✔
476
477	let mut nodes = Vec::new();	1✔
478	pre_order_visit_down(&expr, \|node: &ExprRef\| {	7✔
479	if node.is::<GetItemVTable>() {	7✔
480	nodes.push(node)	2✔
481	}	5✔
482	if let Some(bin) = node.as_opt::<BinaryVTable>() {	7✔
483	if bin.op() == Operator::Eq {	3✔
484	return Ok(TraversalOrder::Skip);	1✔
485	}	2✔
486	}	4✔
487	Ok(TraversalOrder::Continue)	6✔
488	})	7✔
489	.unwrap();	1✔
490
491	let nodes: HashSet<ExprRef> = HashSet::from_iter(nodes.into_iter().cloned());	1✔
492	assert_eq!(nodes, HashSet::from_iter([col("col3"), col("col4")]));	1✔
493	}	1✔
494
495	#[test]
496	fn expr_stop_test() {	1✔
497	let col1: Arc<dyn VortexExpr> = col("col1");	1✔
498	let col2: Arc<dyn VortexExpr> = col("col2");	1✔
499	let expr1 = BinaryExpr::new_expr(col1.clone(), Operator::Eq, col2.clone());	1✔
500	let col3: Arc<dyn VortexExpr> = col("col3");	1✔
501	let col4: Arc<dyn VortexExpr> = col("col4");	1✔
502	let expr2 = BinaryExpr::new_expr(col3.clone(), Operator::NotEq, col4.clone());	1✔
503	let expr = BinaryExpr::new_expr(expr1, Operator::And, expr2);	1✔
504
505	let mut nodes = Vec::new();	1✔
506	pre_order_visit_down(&expr, \|node: &ExprRef\| {	2✔
507	if node.is::<GetItemVTable>() {	2✔
508	nodes.push(node)
509	}	2✔
510	if let Some(bin) = node.as_opt::<BinaryVTable>() {	2✔
511	if bin.op() == Operator::Eq {	2✔
512	return Ok(TraversalOrder::Stop);	1✔
513	}	1✔
514	}
515	Ok(TraversalOrder::Continue)	1✔
516	})	2✔
517	.unwrap();	1✔
518
519	assert!(nodes.is_empty());	1✔
520	}	1✔
521
522	#[test]
523	fn expr_skip_down_visit_up() {	1✔
524	let col = col("col");	1✔
525
526	let mut visitor = SkipDownRewriter;	1✔
527	let result = col.rewrite(&mut visitor).unwrap();	1✔
528
529	assert!(result.changed);	1✔
530	assert!(is_root(&result.value));	1✔
531	}	1✔
532	}

vortex-data / vortex / 16440250234

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous