fix mut slots

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

Author: joseph-isaacs

vortex-array/src/array/mod.rs

@@ -365,6 +365,22 @@ impl dyn DynArray + '_ {
             slots[slot_idx] = Some(replacement.clone());
         })
     }
+
+    /// Take a child out of a slot, setting it to `None`.
+    ///
+    /// Requires unique ownership of the `ArrayRef` (panics if Arc refcount > 1).
+    pub fn take_slot(self: &mut ArrayRef, slot_idx: usize) -> Option<ArrayRef> {
+        let vtable = self.vtable().clone_boxed();
+        vtable.take_slot(self, slot_idx)
+    }
+
+    /// Put a child into a slot, replacing whatever was there.
+    ///
+    /// Requires unique ownership of the `ArrayRef` (panics if Arc refcount > 1).
+    pub fn put_slot(self: &mut ArrayRef, slot_idx: usize, value: ArrayRef) {
+        let vtable = self.vtable().clone_boxed();
+        vtable.put_slot(self, slot_idx, value);
+    }
 }
 
 /// Trait for converting a type into a Vortex [`ArrayRef`].

vortex-array/src/executor.rs

@@ -123,9 +123,9 @@ impl dyn DynArray + '_ {
                         ctx.log(format_args!("-> {}", current));
                         return Ok(current);
                     }
-                    Some((parent, slot_idx, _)) => {
-                        current = parent.with_slot(slot_idx, current)?;
-                        current = current.optimize()?;
+                    Some((mut parent, slot_idx, _)) => {
+                        parent.put_slot(slot_idx, current);
+                        current = parent.optimize()?;
                         continue;
                     }
                 }
@@ -139,9 +139,9 @@ impl dyn DynArray + '_ {
                         ctx.log(format_args!("-> canonical (unmatched) {}", current));
                         return Ok(current);
                     }
-                    Some((parent, slot_idx, _)) => {
-                        current = parent.with_slot(slot_idx, current)?;
-                        current = current.optimize()?;
+                    Some((mut parent, slot_idx, _)) => {
+                        parent.put_slot(slot_idx, current);
+                        current = parent.optimize()?;
                         continue;
                     }
                 }
@@ -159,11 +159,11 @@ impl dyn DynArray + '_ {
 
             // Execute the array itself.
             let result = execute_step(current, ctx)?;
-            let (array, step) = result.into_parts();
+            let (mut array, step) = result.into_parts();
             match step {
                 ExecutionStep::ExecuteSlot(i, done) => {
-                    let child = array.slots()[i]
-                        .clone()
+                    let child = array
+                        .take_slot(i)
                         .vortex_expect("ExecuteSlot index in bounds");
                     ctx.log(format_args!(
                         "ExecuteSlot({i}): pushing {}, focusing on {}",
@@ -338,11 +338,13 @@ impl Executable for ArrayRef {
                 Ok(array)
             }
             ExecutionStep::ExecuteSlot(i, _) => {
-                // For single-step execution, handle ExecuteSlot by executing the slot,
-                // replacing it, and returning the updated array.
-                let child = array.slots()[i].clone().vortex_expect("valid slot index");
+                // Take the child out of the slot, execute it, and put the result back.
+                // This avoids cloning the child or the parent array.
+                let mut array = array;
+                let child = array.take_slot(i).vortex_expect("valid slot index");
                 let executed_child = child.execute::<ArrayRef>(ctx)?;
-                array.with_slot(i, executed_child)
+                array.put_slot(i, executed_child);
+                Ok(array)
             }
         }
     }

vortex-array/src/optimizer/mod.rs

@@ -26,17 +26,16 @@ pub trait ArrayOptimizer {
 
 impl ArrayOptimizer for ArrayRef {
     fn optimize(&self) -> VortexResult<ArrayRef> {
-        Ok(try_optimize(self)?.unwrap_or_else(|| self.clone()))
+        try_optimize(self.clone())
     }
 
     fn optimize_recursive(&self) -> VortexResult<ArrayRef> {
-        Ok(try_optimize_recursive(self)?.unwrap_or_else(|| self.clone()))
+        try_optimize_recursive(self.clone())
     }
 }
 
-fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-    let mut current_array = array.clone();
-    let mut any_optimizations = false;
+fn try_optimize(array: ArrayRef) -> VortexResult<ArrayRef> {
+    let mut current_array = array;
 
     // Apply reduction rules to the current array until no more rules apply.
     let mut loop_counter = 0;
@@ -48,7 +47,6 @@ fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
 
         if let Some(new_array) = current_array.vtable().reduce(&current_array)? {
             current_array = new_array;
-            any_optimizations = true;
             continue;
         }
 
@@ -61,59 +59,28 @@ fn try_optimize(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
                     .vtable()
                     .reduce_parent(child, &current_array, slot_idx)?
             {
-                // If the parent was replaced, then we attempt to reduce it again.
                 current_array = new_array;
-                any_optimizations = true;
-
-                // Continue to the start of the outer loop
                 continue 'outer;
             }
         }
 
-        // No more optimizations can be applied
         break;
     }
 
-    if any_optimizations {
-        Ok(Some(current_array))
-    } else {
-        Ok(None)
-    }
+    Ok(current_array)
 }
 
-fn try_optimize_recursive(array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-    let mut current_array = array.clone();
-    let mut any_optimizations = false;
-
-    if let Some(new_array) = try_optimize(&current_array)? {
-        current_array = new_array;
-        any_optimizations = true;
-    }
+fn try_optimize_recursive(array: ArrayRef) -> VortexResult<ArrayRef> {
+    let mut current_array = try_optimize(array)?;
 
-    // Collect optimized children first, then apply mutations.
-    let slots_snapshot: Vec<_> = current_array.slots().to_vec();
-    let mut optimized_slots: Vec<(usize, ArrayRef)> = Vec::new();
-    for (i, slot) in slots_snapshot.iter().enumerate() {
-        if let Some(child) = slot
-            && let Some(new_child) = try_optimize_recursive(child)?
-        {
-            optimized_slots.push((i, new_child));
+    // Optimize each child slot in-place.
+    let nslots = current_array.slots().len();
+    for i in 0..nslots {
+        if let Some(child) = current_array.take_slot(i) {
+            let optimized = try_optimize_recursive(child)?;
+            current_array.put_slot(i, optimized);
         }
     }
 
-    if !optimized_slots.is_empty() {
-        let vtable = current_array.vtable().clone_boxed();
-        current_array = vtable.with_slots_mut(current_array, &mut |slots| {
-            for (i, new_child) in optimized_slots.drain(..) {
-                slots[i] = Some(new_child);
-            }
-        })?;
-        any_optimizations = true;
-    }
-
-    if any_optimizations {
-        Ok(Some(current_array))
-    } else {
-        Ok(None)
-    }
+    Ok(current_array)
 }

vortex-array/src/vtable/dyn_.rs

@@ -58,6 +58,16 @@ pub trait DynVTable: 'static + Send + Sync + Debug {
         f: &mut dyn FnMut(&mut [Option<ArrayRef>]),
     ) -> VortexResult<ArrayRef>;
 
+    /// Take a child out of a slot, setting it to `None`.
+    ///
+    /// Requires unique ownership of the `ArrayRef` (Arc refcount == 1).
+    fn take_slot(&self, array: &mut ArrayRef, slot_idx: usize) -> Option<ArrayRef>;
+
+    /// Put a child back into a slot.
+    ///
+    /// Requires unique ownership of the `ArrayRef` (Arc refcount == 1).
+    fn put_slot(&self, array: &mut ArrayRef, slot_idx: usize, value: ArrayRef);
+
     /// See [`VTable::reduce`]
     fn reduce(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
 
@@ -131,6 +141,28 @@ impl<V: VTable> DynVTable for V {
         Ok(inner.into_array())
     }
 
+    fn take_slot(&self, array: &mut ArrayRef, slot_idx: usize) -> Option<ArrayRef> {
+        // If we have unique ownership, take in-place. Otherwise, clone the child.
+        if let Some(inner) = downcast_mut::<V>(array) {
+            V::slots_mut(&mut inner.array)[slot_idx].take()
+        } else {
+            array.slots()[slot_idx].clone()
+        }
+    }
+
+    fn put_slot(&self, array: &mut ArrayRef, slot_idx: usize, value: ArrayRef) {
+        // If we have unique ownership, put in-place. Otherwise, rebuild via with_slots_mut.
+        if let Some(inner) = downcast_mut::<V>(array) {
+            V::slots_mut(&mut inner.array)[slot_idx] = Some(value);
+        } else {
+            *array = self
+                .with_slots_mut(array.clone(), &mut |slots| {
+                    slots[slot_idx] = Some(value.clone());
+                })
+                .vortex_expect("put_slot with_slots_mut failed");
+        }
+    }
+
     fn reduce(&self, array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         let Some(reduced) = V::reduce(downcast::<V>(array))? else {
             return Ok(None);
@@ -240,6 +272,16 @@ fn downcast<V: VTable>(array: &ArrayRef) -> &Array<V> {
         .vortex_expect("Failed to downcast array to expected encoding type")
 }
 
+/// Try to get `&mut Array<V>` from an `ArrayRef` if the Arc is uniquely owned.
+///
+/// Returns `None` if the Arc has other references or the type doesn't match.
+fn downcast_mut<V: VTable>(array: &mut ArrayRef) -> Option<&mut Array<V>> {
+    let ptr = Arc::get_mut(array)? as *mut dyn DynArray as *mut Array<V>;
+    // SAFETY: we verified unique ownership via Arc::get_mut, and the vtable on DynVTable<V>
+    // guarantees the concrete type is Array<V>.
+    Some(unsafe { &mut *ptr })
+}
+
 /// Downcast an `ArrayRef` into an `Arc<Array<V>>`.
 fn downcast_owned<V: VTable>(array: ArrayRef) -> Arc<Array<V>> {
     let any_arc = array.as_any_arc();