16460176376

Committed 23 Jul 2025 02:38AM UTC coverage: 81.135% (-0.4%) from 81.5%

Build # 16460176376

Build Type

Pull #3969

github

Committed by

web-flow

Commit Message

Merge f8cef7107 into 2b0334c0b

Pull Request Pull Request #3969: fix: correctly expose gen_test_data cxx ffi

Run Details

42062 of 51842 relevant lines covered (81.13%)

173503.69 hits per line

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

74.93

/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 std::sync::Arc;

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> {
    /// Value that was being rewritten.
    pub value: T,
    /// Controls the flow of rewriting, see [`TraversalOrder`] for more details.
    pub order: TraversalOrder,
    /// Was the value changed during rewriting.
    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
    }

    /// Apply a function to `value`, changing it without changing the `changed` field.
    pub fn map<O, F: FnOnce(T) -> O>(self, f: F) -> Transformed<O> {
        Transformed {
            value: f(self.value),
            order: self.order,
            changed: self.changed,
        }
    }
}

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 Box<C> {
    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
        &'a self,
        f: F,
    ) -> VortexResult<TraversalOrder> {
        self.as_ref().apply_elements(f)
    }

    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
        self,
        f: F,
    ) -> VortexResult<Transformed<Box<C>>> {
        Ok((*self).map_elements(f)?.map(Box::new))
    }
}

impl<'a, T, C> NodeContainer<'a, T> for Arc<C>
where
    T: 'a,
    C: NodeContainer<'a, T> + Clone,
{
    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
        &'a self,
        f: F,
    ) -> VortexResult<TraversalOrder> {
        self.as_ref().apply_elements(f)
    }

    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
        self,
        f: F,
    ) -> VortexResult<Transformed<Arc<C>>> {
        Ok(Arc::unwrap_or_clone(self).map_elements(f)?.map(Arc::new))
    }
}

impl<'a, T: 'a, C: NodeContainer<'a, T>> NodeContainer<'a, T> for [C; 2] {
    fn apply_elements<F: FnMut(&'a T) -> VortexResult<TraversalOrder>>(
        &'a self,
        mut f: F,
    ) -> VortexResult<TraversalOrder> {
        let [lhs, rhs] = self;
        match lhs.apply_elements(&mut f)? {
            TraversalOrder::Skip | TraversalOrder::Continue => rhs.apply_elements(&mut f),
            TraversalOrder::Stop => Ok(TraversalOrder::Stop),
        }
    }

    fn map_elements<F: FnMut(T) -> VortexResult<Transformed<T>>>(
        self,
        mut f: F,
    ) -> VortexResult<Transformed<[C; 2]>> {
        let [lhs, rhs] = self;
        let transformed = lhs.map_elements(&mut f)?;
        match transformed.order {
            TraversalOrder::Skip | TraversalOrder::Continue => {
                let mut t = rhs.map_elements(&mut f)?;
                t.changed |= transformed.changed;
                Ok(t.map(|new_lhs| [new_lhs, transformed.value]))
            }
            TraversalOrder::Stop => Ok(transformed.map(|new_lhs| [new_lhs, rhs])),
        }
    }
}

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

vortex-data / vortex / 16460176376

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