Use Fn for out-of-place lane kernels

vortex-array/src/arrays/decimal/compute/cast.rs

@@ -235,7 +235,7 @@ where
 fn cast_decimal_buffer<F, T>(
     values: &[F],
     valid_values: &Mask,
-    mut cast: impl FnMut(F) -> Option<T>,
+    cast: impl Fn(F) -> Option<T>,
 ) -> Result<Buffer<T>, usize>
 where
     F: NativeDecimalType,
@@ -244,14 +244,14 @@ where
     let mut buffer = BufferMut::<T>::with_capacity(values.len());
     match valid_values {
         Mask::AllTrue(_) => {
-            values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &mut cast)?;
+            values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?;
         }
         Mask::AllFalse(_) => return Ok(BufferMut::<T>::zeroed(values.len()).freeze()),
         Mask::Values(mask) => {
             values.try_map_masked_into(
                 mask.bit_buffer(),
                 &mut buffer.spare_capacity_mut()[..values.len()],
-                &mut cast,
+                &cast,
             )?;
         }
     }

vortex-array/src/arrays/primitive/compute/cast.rs

@@ -316,7 +316,7 @@ fn cast_integer_values_to_decimal_buffer<S, T>(
     values: &[S],
     decimal_dtype: DecimalDType,
     valid_values: &Mask,
-    cast: impl FnMut(S) -> Option<T>,
+    cast: impl Fn(S) -> Option<T>,
 ) -> VortexResult<Buffer<T>>
 where
     S: IntegerPType + ToI256,
@@ -402,7 +402,7 @@ fn decimal_value_fits_precision<T: NativeDecimalType>(
 fn cast_primitive_to_decimal_buffer<S, T>(
     values: &[S],
     valid_values: &Mask,
-    mut cast: impl FnMut(S) -> Option<T>,
+    cast: impl Fn(S) -> Option<T>,
 ) -> Result<Buffer<T>, usize>
 where
     S: NativePType,
@@ -415,13 +415,13 @@ where
     let mut buffer = BufferMut::<T>::with_capacity(values.len());
     match valid_values {
         Mask::AllTrue(_) => {
-            values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &mut cast)?;
+            values.try_map_into(&mut buffer.spare_capacity_mut()[..values.len()], &cast)?;
         }
         Mask::Values(mask) => {
             values.try_map_masked_into(
                 mask.bit_buffer(),
                 &mut buffer.spare_capacity_mut()[..values.len()],
-                &mut cast,
+                &cast,
             )?;
         }
         Mask::AllFalse(_) => unreachable!("all-null values are handled before allocating"),

vortex-array/src/scalar_fn/fns/binary/compare/mod.rs

@@ -293,7 +293,7 @@ pub(super) fn bit_buffer_from_words(words: BufferMut<u64>, len: usize) -> BitBuf
 pub(super) fn collect_zip_bits<T: Copy>(
     lhs: &[T],
     rhs: &[T],
-    mut f: impl FnMut(T, T) -> bool,
+    f: impl Fn(T, T) -> bool,
 ) -> BitBuffer {
     let len = lhs.len();
     let mut words = BufferMut::<u64>::zeroed(len.div_ceil(64));
@@ -302,7 +302,7 @@ pub(super) fn collect_zip_bits<T: Copy>(
 }
 
 /// Bit-pack the predicate `f(values[i])` over a slice into a [`BitBuffer`].
-pub(super) fn collect_bits<T: Copy>(values: &[T], f: impl FnMut(T) -> bool) -> BitBuffer {
+pub(super) fn collect_bits<T: Copy>(values: &[T], f: impl Fn(T) -> bool) -> BitBuffer {
     let len = values.len();
     let mut words = BufferMut::<u64>::zeroed(len.div_ceil(64));
     values.map_bits_into(words.as_mut_slice(), f);

vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs

@@ -49,7 +49,7 @@ pub(super) fn checked_lanes<S, T, Apply>(
 where
     S: IndexedSource,
     T: Copy + Default,
-    Apply: FnMut(S::Item) -> Option<T>,
+    Apply: Fn(S::Item) -> Option<T>,
 {
     let len = source.len();
     debug_assert_eq!(len, valid_rows.len());
@@ -89,13 +89,13 @@ where
 pub(super) fn checked_apply_lanes<S, T, Fail, Apply>(
     source: S,
     valid_rows: &Mask,
-    mut apply: Apply,
+    apply: Apply,
 ) -> Result<Buffer<T>, usize>
 where
     S: IndexedSource + Copy,
     T: Copy + Default,
     Fail: Failure,
-    Apply: FnMut(S::Item) -> (T, Fail),
+    Apply: Fn(S::Item) -> (T, Fail),
 {
     let len = source.len();
     debug_assert_eq!(len, valid_rows.len());
@@ -109,14 +109,14 @@ where
     let mut values = BufferMut::<T>::with_capacity(len);
 
     let out = &mut values.spare_capacity_mut()[..len];
-    if source.map_checked_into(out, &mut apply) != Fail::default() {
-        let mut checked = |item: S::Item| {
+    if source.map_checked_into(out, &apply) != Fail::default() {
+        let checked = |item: S::Item| {
             let (value, failure) = apply(item);
             (failure == Fail::default()).then_some(value)
         };
         match valid_bits {
-            None => source.try_map_into(out, &mut checked)?,
-            Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?,
+            None => source.try_map_into(out, checked)?,
+            Some(valid_bits) => source.try_map_masked_into(valid_bits, out, checked)?,
         }
     }
 

vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs

@@ -211,7 +211,7 @@ where
 fn checked_op_lanes<S, T, Op>(
     source: S,
     valid_rows: &Mask,
-    mut to_operands: impl FnMut(S::Item) -> (T, T),
+    to_operands: impl Fn(S::Item) -> (T, T),
 ) -> Result<Buffer<T>, usize>
 where
     S: IndexedSource + Copy,

vortex-compute/src/lane_kernels/map_into.rs

@@ -18,6 +18,9 @@ use crate::lane_kernels::source::IndexedSource;
 /// `impl<S: IndexedSource> IndexedSourceExt for S` below. Bring the trait into
 /// scope (`use vortex_compute::lane_kernels::IndexedSourceExt;`) to call
 /// them with method syntax: `values.try_map_masked_into(&mask, &mut out, f)`.
+///
+/// Callbacks implement [`Fn`] because each lane must be independent. A callback that mutates
+/// captured state introduces a loop-carried dependency that can prevent vectorization.
 pub trait IndexedSourceExt: IndexedSource + Sized {
     /// Fallible map with mask-aware error attribution. `f` returns `Option<R>`;
     /// `None` indicates a per-lane failure (e.g. range overflow on a narrowing cast).
@@ -28,7 +31,7 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
     /// `is_none()` flags are bit-packed into a `u64` at the lane's position, then
     /// AND-combined with the chunk's validity bitmap — null-lane bits vanish.
     ///
-    /// The closure shape is the same as [`try_map_into`] (`FnMut(Item) -> Option<R>`);
+    /// The closure shape is the same as [`try_map_into`] (`Fn(Item) -> Option<R>`);
     /// the mask parameter is what makes this kernel mask-aware. Callers that need to
     /// distinguish null lanes inside the closure (e.g. to short-circuit an expensive
     /// computation) should construct their own per-lane validity check externally; for
@@ -48,25 +51,25 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
         self,
         mask: &BitBuffer,
         out: &mut [MaybeUninit<R>],
-        mut f: F,
+        f: F,
     ) -> Result<(), usize>
     where
         R: Copy + Default,
-        F: FnMut(Self::Item) -> Option<R>,
+        F: Fn(Self::Item) -> Option<R>,
     {
         #[inline(always)]
         fn chunk<S, R, F>(
             values: &S,
             out: &mut [MaybeUninit<R>],
-            f: &mut F,
+            f: &F,
             src_chunk: u64,
             base: usize,
             count: usize,
         ) -> Option<usize>
         where
             S: IndexedSource,
             R: Copy + Default,
-            F: FnMut(S::Item) -> Option<R>,
+            F: Fn(S::Item) -> Option<R>,
         {
             let mut fail_bits: u64 = 0;
             for bit_idx in 0..count {
@@ -92,15 +95,15 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
         let remainder = len % 64;
 
         for (chunk_idx, src_chunk) in chunks.iter().enumerate() {
-            if let Some(idx) = chunk(&values, out, &mut f, src_chunk, chunk_idx * 64, 64) {
+            if let Some(idx) = chunk(&values, out, &f, src_chunk, chunk_idx * 64, 64) {
                 return Err(idx);
             }
         }
         if remainder != 0
             && let Some(idx) = chunk(
                 &values,
                 out,
-                &mut f,
+                &f,
                 chunks.remainder_bits(),
                 chunks_count * 64,
                 remainder,
@@ -119,20 +122,15 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
     ///
     /// Panics if `out.len() != self.len()`.
     #[inline]
-    fn map_into<R, F>(self, out: &mut [MaybeUninit<R>], mut f: F)
+    fn map_into<R, F>(self, out: &mut [MaybeUninit<R>], f: F)
     where
-        F: FnMut(Self::Item) -> R,
+        F: Fn(Self::Item) -> R,
     {
         #[inline(always)]
-        fn chunk<S, R, F>(
-            values: &S,
-            out: &mut [MaybeUninit<R>],
-            f: &mut F,
-            base: usize,
-            count: usize,
-        ) where
+        fn chunk<S, R, F>(values: &S, out: &mut [MaybeUninit<R>], f: &F, base: usize, count: usize)
+        where
             S: IndexedSource,
-            F: FnMut(S::Item) -> R,
+            F: Fn(S::Item) -> R,
         {
             for bit_idx in 0..count {
                 let idx = base + bit_idx;
@@ -150,10 +148,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
         let remainder = len % CHUNK_LEN;
 
         for chunk_idx in 0..chunks_count {
-            chunk(&values, out, &mut f, chunk_idx * CHUNK_LEN, CHUNK_LEN);
+            chunk(&values, out, &f, chunk_idx * CHUNK_LEN, CHUNK_LEN);
         }
         if remainder != 0 {
-            chunk(&values, out, &mut f, chunks_count * CHUNK_LEN, remainder);
+            chunk(&values, out, &f, chunks_count * CHUNK_LEN, remainder);
         }
     }
 
@@ -179,15 +177,15 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
     ///
     /// Panics if `words.len() < self.len().div_ceil(64)`.
     #[inline]
-    fn map_bits_into<F>(self, words: &mut [u64], mut f: F)
+    fn map_bits_into<F>(self, words: &mut [u64], f: F)
     where
-        F: FnMut(Self::Item) -> bool,
+        F: Fn(Self::Item) -> bool,
     {
         #[inline(always)]
-        fn chunk<S, F>(values: &S, f: &mut F, base: usize, count: usize) -> u64
+        fn chunk<S, F>(values: &S, f: &F, base: usize, count: usize) -> u64
         where
             S: IndexedSource,
-            F: FnMut(S::Item) -> bool,
+            F: Fn(S::Item) -> bool,
         {
             let mut packed: u64 = 0;
             for bit_idx in 0..count {
@@ -211,10 +209,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
         let remainder = len % 64;
 
         for word_idx in 0..full {
-            words[word_idx] = chunk(&values, &mut f, word_idx * 64, 64);
+            words[word_idx] = chunk(&values, &f, word_idx * 64, 64);
         }
         if remainder != 0 {
-            words[full] = chunk(&values, &mut f, full * 64, remainder);
+            words[full] = chunk(&values, &f, full * 64, remainder);
         }
     }
 
@@ -240,10 +238,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
     ///
     /// Panics if `out.len() != self.len()`.
     #[inline]
-    fn map_checked_into<R, Fail, Apply>(self, out: &mut [MaybeUninit<R>], mut apply: Apply) -> Fail
+    fn map_checked_into<R, Fail, Apply>(self, out: &mut [MaybeUninit<R>], apply: Apply) -> Fail
     where
         Fail: Copy + Default + BitOrAssign,
-        Apply: FnMut(Self::Item) -> (R, Fail),
+        Apply: Fn(Self::Item) -> (R, Fail),
     {
         const {
             assert!(
@@ -293,10 +291,10 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
     ///
     /// Panics if `out.len() != self.len()`.
     #[inline]
-    fn try_map_into<R, F>(self, out: &mut [MaybeUninit<R>], mut f: F) -> Result<(), usize>
+    fn try_map_into<R, F>(self, out: &mut [MaybeUninit<R>], f: F) -> Result<(), usize>
     where
         R: Copy + Default,
-        F: FnMut(Self::Item) -> Option<R>,
+        F: Fn(Self::Item) -> Option<R>,
     {
         /// Returns `true` if any lane in `[base, base+count)` failed (OR-reduced);
         /// the cold attribution path is called at the kernel level so it can be
@@ -305,14 +303,14 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
         fn chunk<S, R, F>(
             values: &S,
             out: &mut [MaybeUninit<R>],
-            f: &mut F,
+            f: &F,
             base: usize,
             count: usize,
         ) -> bool
         where
             S: IndexedSource,
             R: Copy + Default,
-            F: FnMut(S::Item) -> Option<R>,
+            F: Fn(S::Item) -> Option<R>,
         {
             let mut fail_acc: u64 = 0;
             for bit_idx in 0..count {
@@ -336,14 +334,14 @@ pub trait IndexedSourceExt: IndexedSource + Sized {
 
         for chunk_idx in 0..chunks_count {
             let base = chunk_idx * CHUNK_LEN;
-            if chunk(&values, out, &mut f, base, CHUNK_LEN) {
-                return Err(attribute_failure_no_mask(&values, base, CHUNK_LEN, &mut f));
+            if chunk(&values, out, &f, base, CHUNK_LEN) {
+                return Err(attribute_failure_no_mask(&values, base, CHUNK_LEN, &f));
             }
         }
         if remainder != 0 {
             let base = chunks_count * CHUNK_LEN;
-            if chunk(&values, out, &mut f, base, remainder) {
-                return Err(attribute_failure_no_mask(&values, base, remainder, &mut f));
+            if chunk(&values, out, &f, base, remainder) {
+                return Err(attribute_failure_no_mask(&values, base, remainder, &f));
             }
         }
         Ok(())
@@ -363,7 +361,7 @@ fn cold_scan<S>(
     values: &S,
     base: usize,
     chunk_len: usize,
-    mut lane_fails: impl FnMut(usize /* bit_idx */, S::Item) -> bool,
+    lane_fails: impl Fn(usize /* bit_idx */, S::Item) -> bool,
 ) -> usize
 where
     S: IndexedSource,
@@ -382,10 +380,10 @@ where
 /// Cold attribution for the no-mask variant. Replays `f` over the chunk to find
 /// the first lane that returns `None`.
 #[inline]
-fn attribute_failure_no_mask<S, R, F>(values: &S, base: usize, chunk_len: usize, f: &mut F) -> usize
+fn attribute_failure_no_mask<S, R, F>(values: &S, base: usize, chunk_len: usize, f: &F) -> usize
 where
     S: IndexedSource,
-    F: FnMut(S::Item) -> Option<R>,
+    F: Fn(S::Item) -> Option<R>,
 {
     cold_scan(values, base, chunk_len, |_bit_idx, val| f(val).is_none())
 }