Minor touchups to vortex-expr (#3466)

There's a lot more work to do there, but I think that's a start.

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

vortex-expr/src/transform/partition.rs

@@ -157,7 +157,7 @@ impl<'a> StructFieldExpressionSplitter<'a> {
         let ctx = ScopeDType::new(dtype.clone());
 
         Ok(PartitionedExpr {
-            root: simplify_typed(split.result, &ctx)?,
+            root: simplify_typed(split.into_inner(), &ctx)?,
             partitions: partitions.into_boxed_slice(),
             partition_names: partition_names.into(),
             partition_dtypes: partition_dtypes.into_boxed_slice(),
@@ -193,7 +193,7 @@ impl FolderMut for StructFieldExpressionSplitter<'_> {
             let replaced = node
                 .clone()
                 .transform(&mut ScopeStepIntoFieldExpr(field_name.clone()))?;
-            sub_exprs.push(replaced.result);
+            sub_exprs.push(replaced.into_inner());
 
             let access = get_item(
                 Self::field_idx_name(field_name, idx),

vortex-expr/src/transform/remove_merge.rs

@@ -6,7 +6,7 @@ use crate::{ExprRef, Merge, ScopeDType, VortexExpr, get_item, pack};
 /// Replaces [Merge] with combination of [GetItem] and [Pack] expressions.
 pub(crate) fn remove_merge(e: ExprRef, ctx: &ScopeDType) -> VortexResult<ExprRef> {
     let mut transform = RemoveMergeTransform { ctx };
-    e.transform(&mut transform).map(|e| e.result)
+    e.transform(&mut transform).map(|e| e.into_inner())
 }
 
 struct RemoveMergeTransform<'a> {

vortex-expr/src/transform/remove_select.rs

@@ -6,7 +6,7 @@ use crate::{ExprRef, ScopeDType, Select, get_item, pack};
 /// Replaces [Select] with combination of [GetItem] and [Pack] expressions.
 pub(crate) fn remove_select(e: ExprRef, ctx: &ScopeDType) -> VortexResult<ExprRef> {
     let mut transform = RemoveSelectTransform { ctx };
-    e.transform(&mut transform).map(|e| e.result)
+    e.transform(&mut transform).map(|e| e.into_inner())
 }
 
 struct RemoveSelectTransform<'a> {

vortex-expr/src/transform/simplify.rs

@@ -10,7 +10,7 @@ use crate::{ExprRef, GetItem, Pack};
 /// If the scope dtype is known, see `simplify_typed` for a simplifier which uses dtype.
 pub fn simplify(e: ExprRef) -> VortexResult<ExprRef> {
     let mut folder = Simplify;
-    let e = e.transform(&mut folder).map(|e| e.result)?;
+    let e = e.transform(&mut folder).map(|e| e.into_inner())?;
     Ok(find_between(e.clone()))
 }
 

vortex-expr/src/traversal/mod.rs

@@ -15,23 +15,22 @@ use crate::ExprRef;
 /// - `Skip`: Skip visiting the children of the current node.
 /// - `Stop`: Stop visiting any more nodes in the traversal.
 /// - `Continue`: Continue with the traversal as expected.
-
 #[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 this does nothing (for now).
+    /// In a top-down traversal, skip visiting the children of the current node.
+    /// In the bottom-up phase of the traversal this does nothing (for now).
     Skip,
 
-    // Stop visiting any more nodes in the traversal.
+    /// Stop visiting any more nodes in the traversal.
     Stop,
 
-    // Continue with the traversal as expected.
+    /// Continue with the traversal as expected.
     Continue,
 }
 
 #[derive(Debug, Clone)]
 pub struct TransformResult<T> {
-    pub result: T,
+    result: T,
     order: TraversalOrder,
     changed: bool,
 }
@@ -52,6 +51,10 @@ impl<T> TransformResult<T> {
             changed: false,
         }
     }
+
+    pub fn into_inner(self) -> T {
+        self.result
+    }
 }
 
 pub trait NodeVisitor<'a> {
@@ -169,7 +172,7 @@ pub trait Node: Sized {
 }
 
 impl Node for ExprRef {
-    // A pre-order traversal.
+    /// A pre-order traversal.
     fn accept<'a, V: NodeVisitor<'a, NodeTy = ExprRef>>(
         &'a self,
         visitor: &mut V,
@@ -216,7 +219,7 @@ impl Node for ExprRef {
         visitor.visit_up(self, context, children)
     }
 
-    // A post-order transform, with an option to ignore sub-tress (using visit_down).
+    /// A post-order transform, with an option to ignore sub-tress (using visit_down).
     fn transform<V: MutNodeVisitor<NodeTy = Self>>(
         self,
         visitor: &mut V,