update

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

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

@@ -5,8 +5,7 @@ use vortex_error::VortexResult;
 
 use crate::expr::exprs::get_item::GetItem;
 use crate::expr::exprs::pack::Pack;
-use crate::expr::transform::ReduceRule;
-use crate::expr::transform::rules::RewriteContext;
+use crate::expr::transform::rules::{ReduceRule, RuleContext};
 use crate::expr::{Expression, ExpressionView};
 
 /// Rewrite rule: `pack(l_1: e_1, ..., l_i: e_i, ..., l_n: e_n).get_item(l_i) = e_i`
@@ -15,11 +14,11 @@ use crate::expr::{Expression, ExpressionView};
 /// expression instead of materializing the pack.
 pub struct PackGetItemRule;
 
-impl ReduceRule<GetItem, &dyn RewriteContext> for PackGetItemRule {
+impl ReduceRule<GetItem, RuleContext> for PackGetItemRule {
     fn reduce(
         &self,
         get_item: &ExpressionView<GetItem>,
-        _ctx: &dyn RewriteContext,
+        _ctx: &RuleContext,
     ) -> VortexResult<Option<Expression>> {
         if let Some(pack) = get_item.child(0).as_opt::<Pack>() {
             let field_expr = pack.field(get_item.data())?;
@@ -41,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::rules::RootRewriteContext;
-    use crate::expr::transform::{ReduceRule, simplify_typed};
+    use crate::expr::transform::rules::{ReduceRule, RuleContext};
+    use crate::expr::transform::simplify_typed;
 
     #[test]
     fn test_pack_get_item_rule() {
@@ -52,12 +51,8 @@ mod tests {
         let pack_expr = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
         let get_item_expr = get_item("b", pack_expr);
 
-        // Create a dummy context
-        let dtype = DType::Primitive(PType::I32, NonNullable);
-        let ctx = RootRewriteContext { dtype: &dtype };
-
         let get_item_view = get_item_expr.as_::<GetItem>();
-        let result = rule.reduce(&get_item_view, &ctx).unwrap();
+        let result = rule.reduce(&get_item_view, &RuleContext).unwrap();
 
         assert!(result.is_some());
         assert_eq!(&result.unwrap(), &lit(2));
@@ -71,11 +66,8 @@ mod tests {
         let lit_expr = lit(42);
         let get_item_expr = get_item("x", lit_expr);
 
-        let dtype = DType::Primitive(PType::I32, NonNullable);
-        let ctx = RootRewriteContext { dtype: &dtype };
-
         let get_item_view = get_item_expr.as_::<GetItem>();
-        let result = rule.reduce(&get_item_view, &ctx).unwrap();
+        let result = rule.reduce(&get_item_view, &RuleContext).unwrap();
 
         assert!(result.is_none());
     }

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

@@ -8,20 +8,19 @@ use vortex_utils::aliases::hash_set::HashSet;
 use crate::expr::exprs::get_item::get_item;
 use crate::expr::exprs::merge::{DuplicateHandling, Merge};
 use crate::expr::exprs::pack::pack;
-use crate::expr::transform::TypedRewriteContext;
-use crate::expr::transform::rules::ReduceRule;
+use crate::expr::transform::rules::{ReduceRule, TypedRuleContext};
 use crate::expr::{Expression, ExpressionView};
 
 /// Rule that removes Merge expressions by converting them to Pack + GetItem.
 ///
 /// Transforms: `merge([struct1, struct2])` → `pack(field1: get_item("field1", struct1), field2: get_item("field2", struct2), ...)`
 pub struct RemoveMergeRule;
 
-impl ReduceRule<Merge, &dyn TypedRewriteContext> for RemoveMergeRule {
+impl ReduceRule<Merge, TypedRuleContext> for RemoveMergeRule {
     fn reduce(
         &self,
         merge: &ExpressionView<Merge>,
-        ctx: &dyn TypedRewriteContext,
+        ctx: &TypedRuleContext,
     ) -> VortexResult<Option<Expression>> {
         let merge_dtype = merge.return_dtype(ctx.dtype())?;
         let mut names = Vec::with_capacity(merge.children().len() * 2);
@@ -82,8 +81,7 @@ mod tests {
     use crate::expr::exprs::merge::{DuplicateHandling, Merge, merge_opts};
     use crate::expr::exprs::pack::Pack;
     use crate::expr::exprs::root::root;
-    use crate::expr::transform::RootRewriteContext;
-    use crate::expr::transform::rules::ReduceRule;
+    use crate::expr::transform::rules::{ReduceRule, TypedRuleContext};
 
     #[test]
     fn test_remove_merge() {
@@ -100,7 +98,7 @@ mod tests {
             DuplicateHandling::RightMost,
         );
 
-        let ctx = RootRewriteContext { dtype: &dtype };
+        let ctx = TypedRuleContext::new(dtype.clone());
         let rule = RemoveMergeRule;
         let merge_view = e.as_::<Merge>();
         let result = rule.reduce(&merge_view, &ctx).unwrap();

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

@@ -6,20 +6,19 @@ use vortex_error::{VortexResult, vortex_err};
 use crate::expr::exprs::get_item::get_item;
 use crate::expr::exprs::pack::pack;
 use crate::expr::exprs::select::Select;
-use crate::expr::transform::TypedRewriteContext;
-use crate::expr::transform::rules::ReduceRule;
+use crate::expr::transform::rules::{ReduceRule, TypedRuleContext};
 use crate::expr::{Expression, ExpressionView};
 
 /// Rule that removes Select expressions by converting them to Pack + GetItem.
 ///
 /// Transforms: `select(["a", "b"], expr)` → `pack(a: get_item("a", expr), b: get_item("b", expr))`
 pub struct RemoveSelectRule;
 
-impl ReduceRule<Select, &dyn TypedRewriteContext> for RemoveSelectRule {
+impl ReduceRule<Select, TypedRuleContext> for RemoveSelectRule {
     fn reduce(
         &self,
         select: &ExpressionView<Select>,
-        ctx: &dyn TypedRewriteContext,
+        ctx: &TypedRuleContext,
     ) -> VortexResult<Option<Expression>> {
         let child = select.child();
         let child_dtype = child.return_dtype(ctx.dtype())?;
@@ -62,7 +61,7 @@ mod tests {
     use crate::expr::exprs::pack::Pack;
     use crate::expr::exprs::root::root;
     use crate::expr::exprs::select::{Select, select};
-    use crate::expr::transform::rules::{ReduceRule, RootRewriteContext};
+    use crate::expr::transform::rules::{ReduceRule, TypedRuleContext};
 
     #[test]
     fn test_remove_select_rule() {
@@ -73,7 +72,7 @@ mod tests {
         let e = select(["a", "b"], root());
 
         let rule = RemoveSelectRule;
-        let ctx = RootRewriteContext { dtype: &dtype };
+        let ctx = TypedRuleContext::new(dtype.clone());
         let select_view = e.as_::<Select>();
         let result = rule.reduce(&select_view, &ctx).unwrap();
 
@@ -97,7 +96,7 @@ mod tests {
         let e = select_exclude(["c"], root());
 
         let rule = RemoveSelectRule;
-        let ctx = RootRewriteContext { dtype: &dtype };
+        let ctx = TypedRuleContext::new(dtype.clone());
         let select_view = e.as_::<Select>();
         let result = rule.reduce(&select_view, &ctx).unwrap();
 

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

@@ -23,8 +23,7 @@ use crate::expr::exprs::pack::Pack;
 use crate::expr::exprs::root::Root;
 use crate::expr::exprs::select::Select;
 use crate::expr::exprs::select::transform::RemoveSelectRule;
-use crate::expr::transform::TypedRewriteContext;
-use crate::expr::transform::rules::{ParentReduceRule, ReduceRule, RewriteContext};
+use crate::expr::transform::rules::{ParentReduceRule, ReduceRule, RuleContext, TypedRuleContext};
 use crate::expr::{ExprVTable, VTable};
 
 /// Registry of expression vtables.
@@ -57,24 +56,24 @@ impl ExprSession {
         self.registry.register_many(exprs);
     }
 
-    /// Register a generic reduce rule that uses TypedRewriteContext.
+    /// Register a generic reduce rule that uses Typed context.
     /// Use this for rules that need access to dtype information.
     pub fn register_typed_reduce_rule<V, R>(&mut self, vtable: &'static V, rule: R)
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ReduceRule<V, &'a dyn TypedRewriteContext>,
+        R: ReduceRule<V, TypedRuleContext>,
     {
         self.rewrite_rules.register_typed_reduce_rule(vtable, rule);
     }
 
-    /// Register a reduce rule that doesn't use TypedRewriteContext.
+    /// Register a reduce rule that uses Untyped context.
     /// Use this for rules that don't need access to dtype information.
     pub fn register_reduce_rule<V, R>(&mut self, vtable: &'static V, rule: R)
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ReduceRule<V, &'a dyn RewriteContext>,
+        R: ReduceRule<V, RuleContext>,
     {
         self.rewrite_rules.register_reduce_rule(vtable, rule);
     }
@@ -84,7 +83,7 @@ impl ExprSession {
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ParentReduceRule<V, &'a dyn RewriteContext>,
+        R: ParentReduceRule<V, RuleContext>,
     {
         self.rewrite_rules.register_parent_rule(vtable, rule);
     }
@@ -94,7 +93,7 @@ impl ExprSession {
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ParentReduceRule<V, &'a dyn TypedRewriteContext>,
+        R: ParentReduceRule<V, TypedRuleContext>,
     {
         self.rewrite_rules.register_typed_parent_rule(vtable, rule);
     }

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

@@ -8,10 +8,9 @@ use std::sync::Arc;
 use vortex_error::VortexResult;
 use vortex_utils::aliases::hash_map::HashMap;
 
-use crate::expr::transform::rules::{ParentReduceRule, ReduceRule, RewriteContext};
+use crate::expr::transform::rules::{ParentReduceRule, ReduceRule, RuleContext, TypedRuleContext};
 use crate::expr::transform::{
     DynParentReduceRule, DynReduceRule, DynTypedParentReduceRule, DynTypedReduceRule,
-    TypedRewriteContext,
 };
 use crate::expr::{ExprId, Expression, VTable};
 
@@ -33,13 +32,9 @@ impl<V: VTable, R> RuleAdapter<V, R> {
 impl<V, R> DynReduceRule for RuleAdapter<V, R>
 where
     V: VTable,
-    for<'a> R: ReduceRule<V, &'a dyn RewriteContext>,
+    R: ReduceRule<V, RuleContext>,
 {
-    fn reduce(
-        &self,
-        expr: &Expression,
-        ctx: &dyn RewriteContext,
-    ) -> VortexResult<Option<Expression>> {
+    fn reduce(&self, expr: &Expression, ctx: &RuleContext) -> VortexResult<Option<Expression>> {
         let Some(view) = expr.as_opt::<V>() else {
             return Ok(None);
         };
@@ -50,12 +45,12 @@ where
 impl<V, R> DynTypedReduceRule for RuleAdapter<V, R>
 where
     V: VTable,
-    for<'a> R: ReduceRule<V, &'a dyn TypedRewriteContext>,
+    R: ReduceRule<V, TypedRuleContext>,
 {
     fn reduce(
         &self,
         expr: &Expression,
-        ctx: &dyn TypedRewriteContext,
+        ctx: &TypedRuleContext,
     ) -> VortexResult<Option<Expression>> {
         let Some(view) = expr.as_opt::<V>() else {
             return Ok(None);
@@ -67,14 +62,14 @@ where
 impl<V, R> DynParentReduceRule for RuleAdapter<V, R>
 where
     V: VTable,
-    for<'a> R: ParentReduceRule<V, &'a dyn RewriteContext>,
+    R: ParentReduceRule<V, RuleContext>,
 {
     fn reduce_parent(
         &self,
         expr: &Expression,
         parent: &Expression,
         child_idx: usize,
-        ctx: &dyn RewriteContext,
+        ctx: &RuleContext,
     ) -> VortexResult<Option<Expression>> {
         let Some(view) = expr.as_opt::<V>() else {
             return Ok(None);
@@ -86,14 +81,14 @@ where
 impl<V, R> DynTypedParentReduceRule for RuleAdapter<V, R>
 where
     V: VTable,
-    for<'a> R: ParentReduceRule<V, &'a dyn TypedRewriteContext>,
+    R: ParentReduceRule<V, TypedRuleContext>,
 {
     fn reduce_parent(
         &self,
         expr: &Expression,
         parent: &Expression,
         child_idx: usize,
-        ctx: &dyn TypedRewriteContext,
+        ctx: &TypedRuleContext,
     ) -> VortexResult<Option<Expression>> {
         let Some(view) = expr.as_opt::<V>() else {
             return Ok(None);
@@ -137,13 +132,13 @@ impl RewriteRuleRegistry {
         Self::default()
     }
 
-    /// Register a generic reduce rule that uses TypedRewriteContext.
+    /// Register a generic reduce rule that uses Typed context.
     /// Use this for rules that need access to dtype information.
     pub fn register_typed_reduce_rule<V, R>(&mut self, vtable: &'static V, rule: R)
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ReduceRule<V, &'a dyn TypedRewriteContext>,
+        R: ReduceRule<V, TypedRuleContext>,
     {
         let adapter = RuleAdapter::new(rule);
         self.typed_reduce_rules
@@ -152,13 +147,13 @@ impl RewriteRuleRegistry {
             .push(Arc::new(adapter));
     }
 
-    /// Register a generic reduce rule that only uses RewriteContext (non-typed).
+    /// Register a generic reduce rule that only uses Untyped context.
     /// Use this for rules that don't need access to dtype information.
     pub fn register_reduce_rule<V, R>(&mut self, vtable: &'static V, rule: R)
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ReduceRule<V, &'a dyn RewriteContext>,
+        R: ReduceRule<V, RuleContext>,
     {
         let adapter = RuleAdapter::new(rule);
         self.reduce_rules
@@ -171,7 +166,7 @@ impl RewriteRuleRegistry {
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ParentReduceRule<V, &'a dyn RewriteContext>,
+        R: ParentReduceRule<V, RuleContext>,
     {
         let adapter = RuleAdapter::new(rule);
         self.parent_rules
@@ -185,7 +180,7 @@ impl RewriteRuleRegistry {
     where
         V: VTable,
         R: 'static,
-        for<'a> R: ParentReduceRule<V, &'a dyn TypedRewriteContext>,
+        R: ParentReduceRule<V, TypedRuleContext>,
     {
         let adapter = RuleAdapter::new(rule);
         self.typed_parent_rules

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

@@ -5,12 +5,14 @@
 pub mod annotations;
 pub mod immediate_access;
 pub(crate) mod match_between;
+mod optimizer;
 mod partition;
 mod replace;
 pub mod rules;
 mod simplify;
 mod simplify_typed;
 
+pub use optimizer::*;
 pub use partition::*;
 pub use replace::*;
 pub use rules::*;