feat[expr]: add back the expression folder, implemented over Node (#4230)

Adds back a Expression Folder

---------

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-expr/src/exprs/dynamic.rs

@@ -14,7 +14,7 @@ use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_proto::expr as pb;
 use vortex_scalar::{Scalar, ScalarValue};
 
-use crate::traversal::{Node, NodeVisitor, TraversalOrder};
+use crate::traversal::{NodeExt, NodeVisitor, TraversalOrder};
 use crate::{
     AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, StatsCatalog, VTable, vtable,
 };

vortex-expr/src/lib.rs

@@ -51,7 +51,7 @@ use vortex_utils::aliases::hash_set::HashSet;
 pub use vtable::*;
 
 use crate::dyn_traits::DynEq;
-use crate::traversal::{Node, ReferenceCollector};
+use crate::traversal::{NodeExt, ReferenceCollector};
 
 pub trait IntoExpr {
     /// Convert this type into an expression reference.

vortex-expr/src/transform/annotations.rs

@@ -8,7 +8,7 @@ use vortex_utils::aliases::hash_map::HashMap;
 use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::ExprRef;
-use crate::traversal::{Node, NodeVisitor, TraversalOrder};
+use crate::traversal::{NodeExt, NodeVisitor, TraversalOrder};
 
 pub trait Annotation: Clone + Hash + Eq {}
 

vortex-expr/src/transform/partition.rs

@@ -12,7 +12,7 @@ use crate::transform::annotations::{
     Annotation, AnnotationFn, Annotations, descendent_annotations,
 };
 use crate::transform::simplify_typed::simplify_typed;
-use crate::traversal::{Node, NodeRewriter, Transformed, TraversalOrder};
+use crate::traversal::{NodeExt, NodeRewriter, Transformed, TraversalOrder};
 use crate::{ExprRef, get_item, pack, root};
 
 /// Partition an expression into sub-expressions that are uniquely associated with an annotation.

vortex-expr/src/transform/remove_merge.rs

@@ -3,7 +3,7 @@
 
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 
-use crate::traversal::{Node, Transformed};
+use crate::traversal::{NodeExt, Transformed};
 use crate::{DType, ExprRef, MergeVTable, get_item, pack};
 
 /// Replaces [crate::MergeExpr] with combination of [crate::GetItem] and [crate::Pack] expressions.

vortex-expr/src/transform/remove_select.rs

@@ -3,7 +3,7 @@
 
 use vortex_error::{VortexResult, vortex_err};
 
-use crate::traversal::{Node, Transformed};
+use crate::traversal::{NodeExt, Transformed};
 use crate::{DType, ExprRef, SelectVTable, get_item, pack};
 
 /// Replaces [crate::SelectExpr] with combination of [crate::GetItem] and [crate::Pack] expressions.

vortex-expr/src/transform/replace.rs

@@ -4,7 +4,7 @@
 use vortex_dtype::{Nullability, StructFields};
 use vortex_error::{VortexExpect, VortexResult};
 
-use crate::traversal::{Node, Transformed};
+use crate::traversal::{NodeExt, Transformed};
 use crate::{ExprRef, col, pack, root};
 
 /// Replaces all occurrences of `needle` in the expression `expr` with `replacement`.

vortex-expr/src/transform/simplify.rs

@@ -5,7 +5,7 @@ use vortex_error::VortexResult;
 
 use crate::transform::match_between::find_between;
 // use crate::transform::match_between::find_between;
-use crate::traversal::{Node, Transformed};
+use crate::traversal::{NodeExt, Transformed};
 use crate::{ExprRef, GetItemVTable, PackVTable};
 
 /// Simplifies an expression into an equivalent expression which is faster and easier to analyze.

vortex-expr/src/traversal/fold.rs

@@ -0,0 +1,210 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::traversal::Node;
+
+pub enum FoldDownContext<C, R> {
+    Continue(C),
+    Stop(R),
+    Skip(R),
+}
+
+pub enum FoldDown<R> {
+    Continue,
+    Stop(R),
+    Skip(R),
+}
+
+pub enum FoldUp<R> {
+    Continue(R),
+    Stop(R),
+}
+
+impl<R> FoldUp<R> {
+    pub fn value(self) -> R {
+        match self {
+            Self::Continue(r) => r,
+            Self::Stop(r) => r,
+        }
+    }
+}
+
+pub trait NodeFolderContext {
+    type NodeTy: Node;
+    type Result;
+    type Context;
+
+    /// visit_down is called when a node is first encountered, in a pre-order traversal.
+    /// If the node's children are to be skipped, return Skip.
+    /// If the node should stop traversal, return Stop.
+    /// Otherwise, return Continue.
+    fn visit_down(
+        &mut self,
+        _ctx: &Self::Context,
+        _node: &Self::NodeTy,
+    ) -> VortexResult<FoldDownContext<Self::Context, Self::Result>>;
+
+    /// visit_up is called when a node is last encountered, in a pre-order traversal.
+    /// If the node should stop traversal, return Stop.
+    /// Otherwise, return Continue.
+    fn visit_up(
+        &mut self,
+        _node: Self::NodeTy,
+        _context: &Self::Context,
+        _children: Vec<Self::Result>,
+    ) -> VortexResult<FoldUp<Self::Result>>;
+}
+
+pub trait NodeFolder {
+    type NodeTy: Node;
+    type Result;
+
+    /// visit_down is called when a node is first encountered, in a pre-order traversal.
+    /// If the node's children are to be skipped, return Skip.
+    /// If the node should stop traversal, return Stop.
+    /// Otherwise, return Continue.
+    fn visit_down(&mut self, _node: &Self::NodeTy) -> VortexResult<FoldDown<Self::Result>>;
+
+    /// visit_up is called when a node is last encountered, in a pre-order traversal.
+    /// If the node should stop traversal, return Stop.
+    /// Otherwise, return Continue.
+    fn visit_up(
+        &mut self,
+        _node: Self::NodeTy,
+        _children: Vec<Self::Result>,
+    ) -> VortexResult<FoldUp<Self::Result>>;
+}
+
+pub struct NodeFolderContextWrapper<'a, T>
+where
+    T: NodeFolder,
+{
+    pub inner: &'a mut T,
+}
+
+impl<T: NodeFolder> NodeFolderContext for NodeFolderContextWrapper<'_, T> {
+    type NodeTy = T::NodeTy;
+    type Result = T::Result;
+    type Context = ();
+
+    fn visit_down(
+        &mut self,
+        _ctx: &Self::Context,
+        _node: &Self::NodeTy,
+    ) -> VortexResult<FoldDownContext<Self::Context, Self::Result>> {
+        match self.inner.visit_down(_node)? {
+            FoldDown::Continue => Ok(FoldDownContext::Continue(())),
+            FoldDown::Stop(r) => Ok(FoldDownContext::Stop(r)),
+            FoldDown::Skip(r) => Ok(FoldDownContext::Skip(r)),
+        }
+    }
+
+    fn visit_up(
+        &mut self,
+        _node: Self::NodeTy,
+        _context: &Self::Context,
+        _children: Vec<Self::Result>,
+    ) -> VortexResult<FoldUp<Self::Result>> {
+        self.inner.visit_up(_node, _children)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_error::VortexExpect;
+
+    use super::*;
+    use crate::traversal::NodeExt;
+    use crate::{
+        BinaryVTable, ExprRef, LiteralVTable, Operator, checked_add, gt, lit, vortex_bail,
+    };
+
+    struct AddFold;
+    impl NodeFolder for AddFold {
+        type NodeTy = ExprRef;
+        type Result = i32;
+
+        fn visit_down(&mut self, node: &'_ Self::NodeTy) -> VortexResult<FoldDown<Self::Result>> {
+            if let Some(lit) = node.as_opt::<LiteralVTable>() {
+                let v = lit
+                    .value()
+                    .as_primitive()
+                    .typed_value::<i32>()
+                    .vortex_expect("i32");
+
+                if v == 5 {
+                    return Ok(FoldDown::Stop(5));
+                }
+            }
+
+            if let Some(binary) = node.as_opt::<BinaryVTable>()
+                && binary.op() == Operator::Gt
+            {
+                return Ok(FoldDown::Skip(0));
+            }
+
+            Ok(FoldDown::Continue)
+        }
+
+        fn visit_up(
+            &mut self,
+            node: Self::NodeTy,
+            children: Vec<Self::Result>,
+        ) -> VortexResult<FoldUp<Self::Result>> {
+            if let Some(lit) = node.as_opt::<LiteralVTable>() {
+                let v = lit
+                    .value()
+                    .as_primitive()
+                    .typed_value::<i32>()
+                    .vortex_expect("i32");
+                Ok(FoldUp::Continue(v))
+            } else if let Some(binary) = node.as_opt::<BinaryVTable>() {
+                if binary.op() == Operator::Add {
+                    Ok(FoldUp::Continue(children[0] + children[1]))
+                } else {
+                    vortex_bail!("not a valid operator")
+                }
+            } else {
+                vortex_bail!("not a valid type")
+            }
+        }
+    }
+
+    #[test]
+    fn test_fold() {
+        let expr = checked_add(checked_add(lit(1), lit(2)), lit(3));
+
+        let mut folder = AddFold;
+        let result = expr.fold(&mut folder).unwrap().value();
+        assert_eq!(result, 6);
+    }
+
+    #[test]
+    fn test_stop_value() {
+        let expr = checked_add(checked_add(lit(1), lit(5)), lit(3));
+
+        let mut folder = AddFold;
+        let result = expr.fold(&mut folder).unwrap().value();
+        assert_eq!(result, 5);
+    }
+
+    #[test]
+    fn test_skip_value() {
+        let expr = checked_add(gt(lit(1), lit(2)), lit(3));
+
+        let mut folder = AddFold;
+        let result = expr.fold(&mut folder).unwrap().value();
+        assert_eq!(result, 3);
+    }
+
+    #[test]
+    fn test_control_flow_value() {
+        let expr = checked_add(gt(lit(1), lit(5)), lit(3));
+
+        let mut folder = AddFold;
+        let result = expr.fold(&mut folder).unwrap().value();
+        assert_eq!(result, 3);
+    }
+}