t

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

Author: joseph-isaacs

vortex-array/src/expr/session/rewrite.rs

@@ -7,8 +7,8 @@ use std::sync::Arc;
 use vortex_error::VortexResult;
 use vortex_utils::aliases::hash_map::HashMap;
 
-use crate::expr::transform::TypedRewriteContext;
 use crate::expr::transform::rules::{ParentReduceRule, ReduceRule, RewriteContext};
+use crate::expr::transform::{Context, TypedRewriteContext};
 use crate::expr::{ExprId, Expression, VTable};
 
 /// Type-erased wrapper for ReduceRule that allows dynamic dispatch.
@@ -120,7 +120,11 @@ pub(crate) trait DynParentReduceRule: Send + Sync {
 }
 
 /// Concrete wrapper that implements DynParentReduceRule for a specific VTable type.
-struct ParentReduceRuleAdapter<V: VTable, R: ParentReduceRule<V>> {
+struct ParentReduceRuleAdapter<V: VTable, R: ParentReduceRule<V, C>>
+where
+    V: VTable,
+    for<'a> R: ReduceRule<V, &'a dyn RewriteContext>,
+{
     rule: R,
     _phantom: PhantomData<V>,
 }
@@ -150,7 +154,7 @@ impl<V: VTable, R: ParentReduceRule<V>> DynParentReduceRule for ParentReduceRule
 }
 
 pub(crate) trait DynTypedParentReduceRule: Send + Sync {
-    fn reduce_parent_dyn(
+    fn reduce_parent_dyn_typed(
         &self,
         expr: &Expression,
         parent: &Expression,
@@ -159,12 +163,20 @@ pub(crate) trait DynTypedParentReduceRule: Send + Sync {
     ) -> VortexResult<Option<Expression>>;
 }
 
-struct TypedParentReduceRuleAdapter<V: VTable, R: ParentReduceRule<V>> {
+struct TypedParentReduceRuleAdapter<V: VTable, R>
+where
+    V: VTable,
+    for<'a> R: ParentReduceRule<V, &'a dyn TypedRewriteContext>,
+{
     rule: R,
     _phantom: PhantomData<V>,
 }
 
-impl<V: VTable, R: ParentReduceRule<V>> TypedParentReduceRuleAdapter<V, R> {
+impl<V, R> TypedParentReduceRuleAdapter<V, R>
+where
+    V: VTable,
+    for<'a> R: ParentReduceRule<V, &'a dyn TypedRewriteContext>,
+{
     fn new(rule: R) -> Self {
         Self {
             rule,
@@ -173,10 +185,12 @@ impl<V: VTable, R: ParentReduceRule<V>> TypedParentReduceRuleAdapter<V, R> {
     }
 }
 
-impl<V: VTable, R: ParentReduceRule<V>> DynTypedParentReduceRule
-    for TypedParentReduceRuleAdapter<V, R>
+impl<V, R> DynTypedParentReduceRule for TypedParentReduceRuleAdapter<V, R>
+where
+    V: VTable,
+    for<'a> R: ParentReduceRule<V, &'a dyn TypedRewriteContext>,
 {
-    fn reduce_parent_dyn(
+    fn reduce_parent_dyn_typed(
         &self,
         expr: &Expression,
         parent: &Expression,
@@ -186,7 +200,7 @@ impl<V: VTable, R: ParentReduceRule<V>> DynTypedParentReduceRule
         let Some(view) = expr.as_opt::<V>() else {
             return Ok(None);
         };
-        self.rule.reduce_parent(&view, parent, child_idx, ctx)
+        self.rule.reduce(&view, parent, child_idx, ctx)
     }
 }
 
@@ -203,7 +217,7 @@ pub struct RewriteRuleRegistry {
     /// Parent reduce rules, indexed by expression ID
     parent_rules: HashMap<ExprId, Vec<Arc<dyn DynParentReduceRule>>>,
     /// Parent reduce rules, indexed by expression ID
-    typed_parent_rules: HashMap<ExprId, Vec<Arc<dyn DynParentReduceRule>>>,
+    typed_parent_rules: HashMap<ExprId, Vec<Arc<dyn DynTypedParentReduceRule>>>,
 }
 
 impl std::fmt::Debug for RewriteRuleRegistry {
@@ -253,7 +267,6 @@ impl RewriteRuleRegistry {
             .push(Arc::new(adapter));
     }
 
-    /// Register a parent reduce rule.
     pub fn register_parent_rule<V: VTable, R: ParentReduceRule<V> + 'static>(
         &mut self,
         vtable: &'static V,
@@ -267,29 +280,52 @@ impl RewriteRuleRegistry {
             .push(Arc::new(adapter));
     }
 
+    /// Register a parent reduce rule.
+    pub fn register_typed_parent_rule<V: VTable, R: ParentReduceRule<V> + 'static>(
+        &mut self,
+        vtable: &'static V,
+        rule: R,
+    ) {
+        let id = vtable.id();
+        let adapter = TypedParentReduceRuleAdapter::new(rule);
+        self.parent_rules
+            .entry(id)
+            .or_default()
+            .push(Arc::new(adapter));
+    }
+
     /// Get all typed reduce rules for a given expression ID.
-    pub(crate) fn typed_reduce_rules_for(
-        &self,
-        id: &ExprId,
-    ) -> Option<&[Arc<dyn DynTypedReduceRule>]> {
-        self.typed_reduce_rules.get(id).map(|v| v.as_slice())
+    pub(crate) fn typed_reduce_rules_for(&self, id: &ExprId) -> &[Arc<dyn DynTypedReduceRule>] {
+        self.typed_reduce_rules
+            .get(id)
+            .map(|v| v.as_slice())
+            .unwrap_or_default()
     }
 
     /// Get all untyped reduce rules for a given expression ID.
-    pub(crate) fn reduce_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn DynReduceRule>]> {
-        self.reduce_rules.get(id).map(|v| v.as_slice())
+    pub(crate) fn reduce_rules_for(&self, id: &ExprId) -> &[Arc<dyn DynReduceRule>] {
+        self.reduce_rules
+            .get(id)
+            .map(|v| v.as_slice())
+            .unwrap_or_default()
     }
 
     /// Get all parent reduce rules for a given expression ID.
-    pub(crate) fn parent_rules_for(&self, id: &ExprId) -> Option<&[Arc<dyn DynParentReduceRule>]> {
-        self.parent_rules.get(id).map(|v| v.as_slice())
+    pub(crate) fn parent_rules_for(&self, id: &ExprId) -> &[Arc<dyn DynParentReduceRule>] {
+        self.parent_rules
+            .get(id)
+            .map(|v| v.as_slice())
+            .unwrap_or_default()
     }
 
     /// Get all the typed parent reduce rules for a given expression ID.
     pub(crate) fn typed_parent_rules_for(
         &self,
         id: &ExprId,
-    ) -> Option<&[Arc<dyn DynParentReduceRule>]> {
-        self.typed_parent_rules.get(id).map(|v| v.as_slice())
+    ) -> &[Arc<dyn DynTypedParentReduceRule>] {
+        self.typed_parent_rules
+            .get(id)
+            .map(|v| v.as_slice())
+            .unwrap_or_default()
     }
 }

vortex-array/src/expr/transform/rules.rs

@@ -42,7 +42,7 @@ pub trait ReduceRule<V: VTable, C: Context>: Send + Sync {
 /// # Type Parameters
 /// * `V` - The VTable type this rule applies to. The rule will only be invoked for expressions
 ///   with this vtable type, providing compile-time type safety.
-pub trait ParentReduceRule<V: VTable>: Send + Sync {
+pub trait ParentReduceRule<V: VTable, C: Context>: Send + Sync {
     /// Try to rewrite an expression based on its parent.
     ///
     /// # Arguments

vortex-array/src/expr/transform/simplify.rs

@@ -16,140 +16,45 @@ use crate::expr::traversal::{NodeExt, Transformed};
 pub fn simplify(e: Expression, session: &ExprSession) -> VortexResult<Expression> {
     let ctx = EmptyRewriteContext;
 
-    // First bottom-up (child rules)
+    let e = apply_parent_rules(e, &ctx, session)?;
     let e = apply_child_rules_impl(e, &ctx, session)?;
-
-    let e = apply_parent_rules_impl(e, &ctx, session)?;
-
     let e = find_between(e);
 
     Ok(e)
 }
 
-/// Internal implementation: Apply parent rules in a top-down manner.
-pub(crate) fn apply_parent_rules_impl(
-    expr: Expression,
-    ctx: &dyn RewriteContext,
-    session: &ExprSession,
-) -> VortexResult<Expression> {
-    apply_parent_rules_recursive(expr, None, ctx, session)
-}
-
-/// Recursive helper for applying parent rules.
-///
-/// This applies parent rules bottom-up:
-/// 1. First recursively process all children
-/// 2. Rebuild expression with new children
-/// 3. Apply parent rules to each child with the rebuilt parent
-/// 4. If any child changes, recursively apply again
-fn apply_parent_rules_recursive(
+fn apply_parent_rules(
     expr: Expression,
-    _parent: Option<&Expression>,
-    ctx: &dyn RewriteContext,
-    session: &ExprSession,
-) -> VortexResult<Expression> {
-    // First, recursively process all children bottom-up
-    let mut new_children = Vec::with_capacity(expr.children().len());
-    let mut children_changed = false;
-
-    for child in expr.children().iter() {
-        // Recursively process this child first
-        let new_child = apply_parent_rules_recursive(child.clone(), Some(&expr), ctx, session)?;
-
-        new_children.push(new_child);
-    }
-
-    // Rebuild the expression with new children if any changed
-    let mut expr = if children_changed {
-        expr.with_children(new_children)?
-    } else {
-        expr
-    };
-
-    // Now apply parent rules to each child using the rebuilt parent
-    loop {
-        let mut any_child_changed = false;
-        let mut updated_children = Vec::with_capacity(expr.children().len());
-
-        for (child_idx, child) in expr.children().iter().enumerate() {
-            // Try to apply parent rules to this child given that expr is its parent
-            let new_child =
-                apply_parent_rules_to_child(child.clone(), &expr, child_idx, ctx, session)?;
-
-            if child != &new_child {
-                any_child_changed = true;
-            }
-
-            updated_children.push(new_child);
-        }
-
-        if any_child_changed {
-            expr = expr.with_children(updated_children)?;
-        } else {
-            break;
-        }
-    }
-
-    Ok(expr)
-}
-
-/// Apply parent rules to a child expression given its parent and child index.
-fn apply_parent_rules_to_child(
-    child: Expression,
-    parent: &Expression,
-    child_idx: usize,
     ctx: &dyn RewriteContext,
     session: &ExprSession,
 ) -> VortexResult<Expression> {
-    let child_id = child.id();
-    if let Some(rules) = session.rewrite_rules().parent_rules_for(&child_id) {
-        let mut current = child;
-        for rule in rules {
-            if let Some(new_expr) = rule.reduce_parent_dyn(&current, parent, child_idx, ctx)? {
-                current = new_expr;
+    expr.transform_up(|node| {
+        for (idx, child) in node.children().iter().enumerate() {
+            for rule in session.rewrite_rules().parent_rules_for(&child.id()) {
+                if let Some(new_expr) = rule.reduce_parent_dyn(&child, &node, idx, ctx)? {
+                    return Ok(Transformed::yes(new_expr));
+                }
             }
         }
-        Ok(current)
-    } else {
-        Ok(child)
-    }
+        Ok(Transformed::no(node))
+    })
+    .map(|t| t.into_inner())
 }
 
-/// Internal implementation: Apply child rules in a bottom-up manner with RewriteContext.
 pub(crate) fn apply_child_rules_impl(
     expr: Expression,
     ctx: &dyn RewriteContext,
     session: &ExprSession,
 ) -> VortexResult<Expression> {
-    expr.transform_up(|node| apply_reduce_rules_node(node, ctx, session))
-        .map(|t| t.into_inner())
-}
-
-/// Apply child rules to a single node with RewriteContext.
-fn apply_reduce_rules_node(
-    expr: Expression,
-    ctx: &dyn RewriteContext,
-    session: &ExprSession,
-) -> VortexResult<Transformed<Expression>> {
-    let expr_id = expr.id();
-    let mut current = expr;
-    let mut changed = false;
-
-    // Apply untyped generic reduce rules
-    if let Some(rules) = session.rewrite_rules().reduce_rules_for(&expr_id) {
-        for rule in rules {
-            if let Some(new_expr) = rule.reduce_dyn(&current, ctx)? {
-                current = new_expr;
-                changed = true;
+    expr.transform_down(|node| {
+        for rule in session.rewrite_rules().reduce_rules_for(&node.id()) {
+            if let Some(new_expr) = rule.reduce_dyn(&node, ctx)? {
+                return Ok(Transformed::yes(new_expr));
             }
         }
-    }
-
-    if changed {
-        Ok(Transformed::yes(current))
-    } else {
-        Ok(Transformed::no(current))
-    }
+        Ok(Transformed::no(node))
+    })
+    .map(|t| t.into_inner())
 }
 
 #[cfg(test)]