vortex_array/expr/traversal/
fold.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::expr::traversal::Node;
7
8/// Use to indicate the control flow of the fold on the downwards pass.
9/// In the case of Continue, the context is passed on to the children nodes.
10/// `Stop` ends the fold with a result, while `Skip` skips the current node's children.
11#[derive(Debug)]
12pub enum FoldDownContext<C, R> {
13    Continue(C),
14    Stop(R),
15    Skip(R),
16}
17
18/// Use to indicate the control flow of the fold on the upwards pass.
19/// `Stop` indicates that the fold should stop at the current position and return the result.
20#[derive(Debug)]
21pub enum FoldUp<R> {
22    Continue(R),
23    Stop(R),
24}
25
26impl<R> FoldUp<R> {
27    pub fn value(self) -> R {
28        match self {
29            Self::Continue(r) => r,
30            Self::Stop(r) => r,
31        }
32    }
33}
34
35/// Use to implement the folding a tree like structure in a pre-order traversal.
36///
37/// At each point on the way down, the `visit_down` method is called. If it returns `Skip`,
38/// the children of the current node are skipped. If it returns `Stop`, the fold is stopped.
39/// If it returns `Continue`, the children of the current node are visited.
40///
41/// At each point on the way up, the `visit_up` method is called. If it returns `Stop`,
42/// the fold stops.
43///
44/// On the way up the folded children are passed to the `visit_up` method along with the current node.
45///
46/// Note: this trait is not safe to use for graphs with a cycle.
47pub trait NodeFolderContext {
48    type NodeTy: Node;
49    type Result;
50    type Context;
51
52    /// visit_down is called when a node is first encountered, in a pre-order traversal.
53    /// If the node's children are to be skipped, return Skip.
54    /// If the node should stop traversal, return Stop.
55    /// Otherwise, return Continue.
56    fn visit_down(
57        &mut self,
58        _ctx: &Self::Context,
59        _node: &Self::NodeTy,
60    ) -> VortexResult<FoldDownContext<Self::Context, Self::Result>>;
61
62    /// visit_up is called when a node is last encountered, in a pre-order traversal.
63    /// If the node should stop traversal, return Stop.
64    /// Otherwise, return Continue.
65    fn visit_up(
66        &mut self,
67        _node: Self::NodeTy,
68        _context: &Self::Context,
69        _children: Vec<Self::Result>,
70    ) -> VortexResult<FoldUp<Self::Result>>;
71}
72
73#[cfg(test)]
74mod tests {
75    use vortex_error::VortexExpect;
76    use vortex_error::vortex_bail;
77
78    use super::*;
79    use crate::expr::Expression;
80    use crate::expr::checked_add;
81    use crate::expr::gt;
82    use crate::expr::lit;
83    use crate::expr::traversal::NodeExt;
84    use crate::scalar_fn::fns::binary::Binary;
85    use crate::scalar_fn::fns::literal::Literal;
86    use crate::scalar_fn::fns::operators::Operator;
87
88    struct AddFold;
89    impl NodeFolderContext for AddFold {
90        type NodeTy = Expression;
91        type Result = i32;
92        type Context = ();
93
94        fn visit_down(
95            &mut self,
96            _ctx: &Self::Context,
97            node: &Self::NodeTy,
98        ) -> VortexResult<FoldDownContext<Self::Context, Self::Result>> {
99            if let Some(scalar) = node.as_opt::<Literal>() {
100                let v = scalar
101                    .as_primitive()
102                    .typed_value::<i32>()
103                    .vortex_expect("i32");
104
105                if v == 5 {
106                    return Ok(FoldDownContext::Stop(5));
107                }
108            }
109
110            if let Some(operator) = node.as_opt::<Binary>()
111                && *operator == Operator::Gt
112            {
113                return Ok(FoldDownContext::Skip(0));
114            }
115
116            Ok(FoldDownContext::Continue(()))
117        }
118
119        fn visit_up(
120            &mut self,
121            node: Self::NodeTy,
122            _context: &Self::Context,
123            children: Vec<Self::Result>,
124        ) -> VortexResult<FoldUp<Self::Result>> {
125            if let Some(scalar) = node.as_opt::<Literal>() {
126                let v = scalar
127                    .as_primitive()
128                    .typed_value::<i32>()
129                    .vortex_expect("i32");
130                Ok(FoldUp::Continue(v))
131            } else if let Some(operator) = node.as_opt::<Binary>() {
132                if *operator == Operator::Add {
133                    Ok(FoldUp::Continue(children[0] + children[1]))
134                } else {
135                    vortex_bail!("not a valid operator")
136                }
137            } else {
138                vortex_bail!("not a valid type")
139            }
140        }
141    }
142
143    #[test]
144    fn test_fold() {
145        let expr = checked_add(checked_add(lit(1), lit(2)), lit(3));
146
147        let mut folder = AddFold;
148        let result = expr.fold_context(&(), &mut folder).unwrap().value();
149        assert_eq!(result, 6);
150    }
151
152    #[test]
153    fn test_stop_value() {
154        let expr = checked_add(checked_add(lit(1), lit(5)), lit(3));
155
156        let mut folder = AddFold;
157        let result = expr.fold_context(&(), &mut folder).unwrap().value();
158        assert_eq!(result, 5);
159    }
160
161    #[test]
162    fn test_skip_value() {
163        let expr = checked_add(gt(lit(1), lit(2)), lit(3));
164
165        let mut folder = AddFold;
166        let result = expr.fold_context(&(), &mut folder).unwrap().value();
167        assert_eq!(result, 3);
168    }
169
170    #[test]
171    fn test_control_flow_value() {
172        let expr = checked_add(gt(lit(1), lit(5)), lit(3));
173
174        let mut folder = AddFold;
175        let result = expr.fold_context(&(), &mut folder).unwrap().value();
176        assert_eq!(result, 3);
177    }
178}
vortex_array/expr/traversal/fold.rs

vortex_array/expr/traversal/
fold.rs
