Add expression transformation to replace all occurrences of a node (#3795)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-expr/src/exprs/pack.rs

@@ -199,7 +199,7 @@ impl Display for PackExpr {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
             f,
-            "pack({{{}}}){}",
+            "pack({}){}",
             self.names
                 .iter()
                 .zip(&self.values)

vortex-expr/src/transform/mod.rs

@@ -9,6 +9,7 @@ pub(crate) mod match_between;
 pub mod partition;
 mod remove_merge;
 mod remove_select;
+pub mod replace;
 pub mod simplify;
 pub mod simplify_typed;
 pub mod var_partition;

vortex-expr/src/transform/replace.rs

@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::{VortexExpect, VortexResult};
+
+use crate::traversal::{MutNodeVisitor, Node, TransformResult, TraversalOrder};
+use crate::{ExprRef, VortexExpr};
+
+/// Replaces all occurrences of `needle` in the expression `expr` with `replacement`.
+pub fn replace(expr: ExprRef, needle: &dyn VortexExpr, replacement: ExprRef) -> ExprRef {
+    let mut transform = ReplaceVisitor {
+        needle,
+        replacement,
+    };
+    expr.transform(&mut transform)
+        .vortex_expect("ReplaceVisitor should not fail")
+        .into_inner()
+}
+
+/// A visitor that replaces occurrences of a specific expression (`needle`) with a replacement
+/// expression (`replacement`).
+struct ReplaceVisitor<'a> {
+    needle: &'a dyn VortexExpr,
+    replacement: ExprRef,
+}
+
+impl MutNodeVisitor for ReplaceVisitor<'_> {
+    type NodeTy = ExprRef;
+
+    fn visit_down(&mut self, node: &Self::NodeTy) -> VortexResult<TraversalOrder> {
+        if self.needle.eq(node.as_ref()) {
+            // Short-circuit traversal if the needle is found
+            Ok(TraversalOrder::Skip)
+        } else {
+            Ok(TraversalOrder::Continue)
+        }
+    }
+
+    fn visit_up(&mut self, node: Self::NodeTy) -> VortexResult<TransformResult<Self::NodeTy>> {
+        if self.needle.eq(node.as_ref()) {
+            Ok(TransformResult::yes(self.replacement.clone()))
+        } else {
+            Ok(TransformResult::no(node))
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use vortex_dtype::Nullability::NonNullable;
+
+    use super::replace;
+    use crate::{get_item, lit, pack};
+
+    #[test]
+    fn test_replace_full_tree() {
+        let e = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
+        let needle = get_item("b", pack([("a", lit(1)), ("b", lit(2))], NonNullable));
+        let replacement = lit(42);
+        let replaced_expr = replace(e, needle.as_ref(), replacement.clone());
+        assert_eq!(&replaced_expr, &replacement);
+    }
+
+    #[test]
+    fn test_replace_leaf() {
+        let e = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
+        let needle = lit(2);
+        let replacement = lit(42);
+        let replaced_expr = replace(e, needle.as_ref(), replacement.clone());
+        assert_eq!(replaced_expr.to_string(), "pack(a: 1i32, b: 42i32)");
+    }
+}