refactor: simplify array optimizer

vortex-array/src/optimizer/mod.rs

@@ -22,12 +22,17 @@ use vortex_error::vortex_bail;
 use vortex_session::VortexSession;
 
 use crate::ArrayRef;
+use crate::optimizer::kernels::ArrayKernels;
 use crate::optimizer::kernels::ArrayKernelsExt;
 use crate::trace_op;
 
 pub mod kernels;
 pub mod rules;
 
+/// Last zero-based fixpoint pass attempted before treating continued rewrites as an infinite loop.
+/// Increasing this value permits longer rewrite chains but delays detection of cyclic rules.
+const MAX_OPTIMIZER_REWRITE_PASS: usize = 100;
+
 /// Extension trait for optimizing array trees using reduce/reduce_parent rules.
 pub trait ArrayOptimizer {
     /// Optimize the root array node by running reduce and reduce_parent rules to fixpoint.
@@ -72,12 +77,11 @@ fn try_optimize(
 ) -> VortexResult<Option<ArrayRef>> {
     let mut current_array = array.clone();
     let mut any_optimizations = false;
-    let array_ref = session.map(|s| s.kernels());
+    let session_kernels = session.map(|session| session.kernels());
 
     trace_op!(record_optimize_start(array, session.is_some()));
 
-    // Apply reduction rules to the current array until no more rules apply.
-    for _ in 0..=100 {
+    for _ in 0..=MAX_OPTIMIZER_REWRITE_PASS {
         trace_op!(record_optimize_loop_start(&current_array));
 
         if let Some(new_array) = current_array.reduce()? {
@@ -89,50 +93,29 @@ fn try_optimize(
 
         trace_op!(record_optimize_reduce_none(&current_array));
 
-        // Apply parent reduction rules to each slot in the context of the current array.
-        // Its important to take all slots here, as `current_array` can change inside the loop.
-        let mut parent_reduced = None;
+        // Try children in order; the first parent rewrite restarts the fixpoint loop.
+        let mut reduced_parent = None;
         for (slot_idx, slot) in current_array.slots().iter().enumerate() {
-            let Some(child) = slot else { continue };
-
-            // Session kernels take precedence over the child encoding's static PARENT_RULES.
-            if let Some(array_ref) = &array_ref
-                && let Some(plugins) =
-                    array_ref.find_reduce_parent(current_array.encoding_id(), child.encoding_id())
+            let Some(child) = slot else {
+                continue;
+            };
+
+            // Session kernels take precedence over the child's static parent-reduce rules.
+            if let Some(session_kernels) = &session_kernels
+                && let Some(new_array) =
+                    try_session_parent_reduce(session_kernels, &current_array, child, slot_idx)?
             {
-                #[allow(clippy::unused_enumerate_index)]
-                for (_plugin_idx, plugin) in plugins.as_ref().iter().enumerate() {
-                    if let Some(new_array) = plugin(child, &current_array, slot_idx)? {
-                        trace_op!(record_session_parent_reduce_applied(
-                            &current_array,
-                            child,
-                            slot_idx,
-                            _plugin_idx,
-                            &new_array,
-                        ));
-                        parent_reduced = Some(new_array);
-                        break;
-                    }
-                    trace_op!(record_session_parent_reduce_declined(
-                        &current_array,
-                        child,
-                        slot_idx,
-                        _plugin_idx,
-                    ));
-                }
-                if parent_reduced.is_some() {
-                    break;
-                }
+                reduced_parent = Some(new_array);
+                break;
             }
 
             if let Some(new_array) = child.reduce_parent(&current_array, slot_idx)? {
-                parent_reduced = Some(new_array);
+                reduced_parent = Some(new_array);
                 break;
             }
         }
 
-        if let Some(new_array) = parent_reduced {
-            // If the parent was replaced, then we attempt to reduce it again.
+        if let Some(new_array) = reduced_parent {
             current_array = new_array;
             any_optimizations = true;
             trace_op!(record_optimize_loop_end());
@@ -142,19 +125,51 @@ fn try_optimize(
         trace_op!(record_optimize_parent_reduce_none(&current_array));
         trace_op!(record_optimize_loop_end());
 
-        // No more optimizations can be applied
         trace_op!(record_optimize_done(&current_array, any_optimizations));
 
-        if any_optimizations {
-            return Ok(Some(current_array));
-        } else {
-            return Ok(None);
-        }
+        return Ok(any_optimizations.then_some(current_array));
     }
 
     vortex_bail!("Exceeded maximum optimization iterations (possible infinite loop)");
 }
 
+fn try_session_parent_reduce(
+    kernels: &ArrayKernels,
+    parent: &ArrayRef,
+    child: &ArrayRef,
+    slot_idx: usize,
+) -> VortexResult<Option<ArrayRef>> {
+    let Some(reduce_parent_fns) =
+        kernels.find_reduce_parent(parent.encoding_id(), child.encoding_id())
+    else {
+        return Ok(None);
+    };
+
+    #[allow(clippy::unused_enumerate_index)]
+    for (_kernel_idx, reduce_parent) in reduce_parent_fns.iter().enumerate() {
+        if let Some(new_array) = reduce_parent(child, parent, slot_idx)? {
+            trace_op!(record_session_parent_reduce_applied(
+                parent,
+                child,
+                slot_idx,
+                _kernel_idx,
+                &new_array,
+            ));
+
+            return Ok(Some(new_array));
+        }
+
+        trace_op!(record_session_parent_reduce_declined(
+            parent,
+            child,
+            slot_idx,
+            _kernel_idx,
+        ));
+    }
+
+    Ok(None)
+}
+
 fn try_optimize_recursive(
     array: &ArrayRef,
     session: &VortexSession,