cleanup

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

Author: joseph-isaacs

vortex-array/src/expr/exprs/get_item/transform.rs

@@ -40,8 +40,8 @@ mod tests {
     use crate::expr::exprs::literal::lit;
     use crate::expr::exprs::pack::pack;
     use crate::expr::session::ExprSession;
+    use crate::expr::transform::ExprOptimizer;
     use crate::expr::transform::rules::{ReduceRule, RuleContext};
-    use crate::expr::transform::simplify_typed;
 
     #[test]
     fn test_pack_get_item_rule() {
@@ -82,7 +82,9 @@ mod tests {
 
         let dtype = DType::Primitive(PType::I32, NonNullable);
 
-        let result = simplify_typed(get_z, &dtype, ExprSession::default().rewrite_rules()).unwrap();
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
+        let result = optimizer.optimize_typed(get_z, &dtype).unwrap();
 
         assert_eq!(&result, &lit(4));
     }
@@ -103,8 +105,9 @@ mod tests {
 
         let dtype = DType::Primitive(PType::I32, NonNullable);
 
-        let result =
-            simplify_typed(get_final, &dtype, ExprSession::default().rewrite_rules()).unwrap();
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
+        let result = optimizer.optimize_typed(get_final, &dtype).unwrap();
 
         assert_eq!(&result, &lit(42));
     }
@@ -120,8 +123,9 @@ mod tests {
 
         let dtype = DType::Primitive(PType::I32, NonNullable);
 
-        let result =
-            simplify_typed(get_result, &dtype, ExprSession::default().rewrite_rules()).unwrap();
+        let session = ExprSession::default();
+        let optimizer = ExprOptimizer::new(&session);
+        let result = optimizer.optimize_typed(get_result, &dtype).unwrap();
 
         let expected = checked_add(lit(1), lit(10));
         assert_eq!(&result, &expected);

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

@@ -105,7 +105,7 @@ where
 }
 
 type RuleRegistry<Rule> = DashMap<ExprId, Vec<Arc<Rule>>>;
-type ParentRuleRegistry<Rule> = DashMap<ExprId, Vec<Arc<Rule>>>;
+type ParentRuleRegistry<Rule> = DashMap<(ExprId, ExprId), Vec<Arc<Rule>>>;
 
 /// Inner struct that holds all the rule registries.
 /// Wrapped in a single Arc by RewriteRuleRegistry for efficient cloning.
@@ -218,7 +218,7 @@ impl RewriteRuleRegistry {
         };
         self.inner
             .parent_rules
-            .entry(child_vtable.id())
+            .entry((child_vtable.id(), parent_vtable.id()))
             .or_default()
             .push(Arc::new(adapter));
     }
@@ -235,7 +235,7 @@ impl RewriteRuleRegistry {
             _phantom: PhantomData,
         };
         self.inner
-            .parent_rules
+            .any_parent_rules
             .entry(child_vtable.id())
             .or_default()
             .push(Arc::new(adapter));
@@ -259,7 +259,7 @@ impl RewriteRuleRegistry {
         };
         self.inner
             .typed_parent_rules
-            .entry(child_vtable.id())
+            .entry((child_vtable.id(), parent_vtable.id()))
             .or_default()
             .push(Arc::new(adapter));
     }
@@ -279,7 +279,7 @@ impl RewriteRuleRegistry {
             _phantom: PhantomData,
         };
         self.inner
-            .typed_parent_rules
+            .typed_any_parent_rules
             .entry(child_vtable.id())
             .or_default()
             .push(Arc::new(adapter));
@@ -288,48 +288,84 @@ impl RewriteRuleRegistry {
     /// Execute a callback with all typed reduce rules for a given expression ID.
     pub(crate) fn with_typed_reduce_rules<F, R>(&self, id: &ExprId, f: F) -> R
     where
-        F: FnOnce(&[Arc<dyn DynTypedReduceRule>]) -> R,
+        F: FnOnce(&mut dyn Iterator<Item = &dyn DynTypedReduceRule>) -> R,
     {
-        if let Some(entry) = self.inner.typed_reduce_rules.get(id) {
-            f(entry.value())
-        } else {
-            f(&[])
-        }
+        f(&mut self
+            .inner
+            .typed_reduce_rules
+            .get(id)
+            .iter()
+            .map(|v| v.value())
+            .flatten()
+            .map(|arc| arc.as_ref()))
     }
 
     /// Execute a callback with all untyped reduce rules for a given expression ID.
     pub(crate) fn with_reduce_rules<F, R>(&self, id: &ExprId, f: F) -> R
     where
-        F: FnOnce(&[Arc<dyn DynReduceRule>]) -> R,
+        F: FnOnce(&mut dyn Iterator<Item = &dyn DynReduceRule>) -> R,
     {
-        if let Some(entry) = self.inner.reduce_rules.get(id) {
-            f(entry.value())
-        } else {
-            f(&[])
-        }
+        f(&mut self
+            .inner
+            .reduce_rules
+            .get(id)
+            .iter()
+            .map(|v| v.value())
+            .flatten()
+            .map(|arc| arc.as_ref()))
     }
 
-    /// Execute a callback with all untyped parent reduce rules for a given expression ID.
-    pub(crate) fn with_parent_rules<F, R>(&self, id: &ExprId, f: F) -> R
+    /// Execute a callback with all untyped parent reduce rules for a given child and parent expression ID.
+    ///
+    /// Returns rules from both specific parent rules (if parent_id provided) and "any parent" wildcard rules.
+    pub(crate) fn with_parent_rules<F, R>(
+        &self,
+        child_id: &ExprId,
+        parent_id: Option<&ExprId>,
+        f: F,
+    ) -> R
     where
-        F: FnOnce(&[Arc<dyn DynParentReduceRule>]) -> R,
+        F: FnOnce(&mut dyn Iterator<Item = &dyn DynParentReduceRule>) -> R,
     {
-        if let Some(entry) = self.inner.parent_rules.get(id) {
-            f(entry.value())
-        } else {
-            f(&[])
-        }
+        let specific_entry = parent_id.and_then(|pid| {
+            self.inner
+                .parent_rules
+                .get(&(child_id.clone(), pid.clone()))
+        });
+        let wildcard_entry = self.inner.any_parent_rules.get(child_id);
+
+        f(&mut specific_entry
+            .iter()
+            .map(|v| v.value())
+            .flatten()
+            .chain(wildcard_entry.iter().map(|v| v.value()).flatten())
+            .map(|arc| arc.as_ref()))
     }
 
-    /// Execute a callback with all typed parent reduce rules for a given expression ID.
-    pub(crate) fn with_typed_parent_rules<F, R>(&self, id: &ExprId, f: F) -> R
+    /// Execute a callback with all typed parent reduce rules for a given child and parent expression ID.
+    ///
+    /// Returns rules from both specific parent rules (if parent_id provided) and "any parent" wildcard rules.
+    pub(crate) fn with_typed_parent_rules<F, R>(
+        &self,
+        child_id: &ExprId,
+        parent_id: Option<&ExprId>,
+        f: F,
+    ) -> R
     where
-        F: FnOnce(&[Arc<dyn DynTypedParentReduceRule>]) -> R,
+        F: FnOnce(&mut dyn Iterator<Item = &dyn DynTypedParentReduceRule>) -> R,
     {
-        if let Some(entry) = self.inner.typed_parent_rules.get(id) {
-            f(entry.value())
-        } else {
-            f(&[])
-        }
+        let specific_entry = parent_id.and_then(|pid| {
+            self.inner
+                .typed_parent_rules
+                .get(&(child_id.clone(), pid.clone()))
+        });
+        let wildcard_entry = self.inner.typed_any_parent_rules.get(child_id);
+
+        f(&mut specific_entry
+            .iter()
+            .map(|v| v.value())
+            .flatten()
+            .chain(wildcard_entry.iter().map(|v| v.value()).flatten())
+            .map(|arc| arc.as_ref()))
     }
 }

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

@@ -16,5 +16,3 @@ pub use optimizer::*;
 pub use partition::*;
 pub use replace::*;
 pub use rules::*;
-pub(crate) use simplify::*;
-pub(crate) use simplify_typed::*;

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

@@ -6,27 +6,23 @@ use vortex_error::VortexResult;
 
 use crate::expr::Expression;
 use crate::expr::session::{ExprSession, RewriteRuleRegistry};
-use crate::expr::transform::{simplify, simplify_typed};
+use crate::expr::transform::simplify::simplify;
+use crate::expr::transform::simplify_typed::simplify_typed;
 
 /// A unified optimizer for expressions that can work with or without type information.
 pub struct ExprOptimizer {
     rule_registry: RewriteRuleRegistry,
 }
 
 impl ExprOptimizer {
-    /// Create a new untyped optimizer.
-    ///
-    /// This optimizer will use untyped simplification rules only.
+    /// Creates a new optimizer with the rules in `ExprSession`.
     pub fn new(session: &ExprSession) -> Self {
         Self {
             rule_registry: session.rewrite_rules().clone(),
         }
     }
 
-    /// Optimize the given expression.
-    ///
-    /// If this optimizer was created with a dtype, this will perform typed optimization.
-    /// Otherwise, it will perform untyped optimization.
+    /// Optimize the given expression without a dtype.
     pub fn optimize(&self, expr: Expression) -> VortexResult<Expression> {
         simplify(expr, &self.rule_registry)
     }

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

@@ -13,7 +13,7 @@ use crate::expr::traversal::{NodeExt, Transformed};
 ///
 /// This applies only untyped rewrite rules registered in the default session.
 /// If the scope dtype is known, see `simplify_typed` for a simplifier which uses dtype.
-pub(crate) fn simplify(
+pub(super) fn simplify(
     e: Expression,
     rule_registry: &RewriteRuleRegistry,
 ) -> VortexResult<Expression> {
@@ -35,6 +35,7 @@ fn apply_parent_rules(
         for (idx, child) in node.children().iter().enumerate() {
             let result = rule_registry.with_parent_rules(
                 &child.id(),
+                Some(&node.id()),
                 |rules| -> VortexResult<Option<Expression>> {
                     for rule in rules {
                         if let Some(new_expr) = rule.reduce_parent(child, &node, idx, ctx)? {
@@ -53,7 +54,7 @@ fn apply_parent_rules(
     .map(|t| t.into_inner())
 }
 
-pub(crate) fn apply_child_rules_impl(
+fn apply_child_rules_impl(
     expr: Expression,
     ctx: &RuleContext,
     rule_registry: &RewriteRuleRegistry,

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

@@ -15,7 +15,7 @@ use crate::expr::traversal::{NodeExt, Transformed};
 ///
 /// NOTE: After typed simplification, returned expressions is "bound" to the scope DType.
 ///     Applying the returned expression to a different DType may produce wrong results.
-pub(crate) fn simplify_typed(
+pub(super) fn simplify_typed(
     expr: Expression,
     dtype: &DType,
     rule_registry: &RewriteRuleRegistry,
@@ -88,6 +88,7 @@ fn apply_parent_rules_impl_typed(
         for (idx, child) in node.children().iter().enumerate() {
             let result = rule_registry.with_typed_parent_rules(
                 &child.id(),
+                Some(&node.id()),
                 |rules| -> VortexResult<Option<Expression>> {
                     for rule in rules {
                         if let Some(new_expr) = rule.reduce_parent(child, &node, idx, ctx)? {
@@ -105,6 +106,7 @@ fn apply_parent_rules_impl_typed(
             let untyped_ctx: RuleContext = ctx.into();
             let result = rule_registry.with_parent_rules(
                 &child.id(),
+                Some(&node.id()),
                 |rules| -> VortexResult<Option<Expression>> {
                     for rule in rules {
                         if let Some(new_expr) =