Lengthen life in map_inplace and unify impl with fold

Author: jturner314

src/dimension/mod.rs

@@ -678,6 +678,36 @@ where
     }
 }
 
+/// Move the axis which has the smallest absolute stride and a length
+/// greater than one to be the last axis.
+pub fn move_min_stride_axis_to_last<D>(dim: &mut D, strides: &mut D)
+where
+    D: Dimension,
+{
+    debug_assert_eq!(dim.ndim(), strides.ndim());
+    match dim.ndim() {
+        0 | 1 => {}
+        2 => {
+            if dim[1] <= 1
+                || dim[0] > 1 && (strides[0] as isize).abs() < (strides[1] as isize).abs()
+            {
+                dim.slice_mut().swap(0, 1);
+                strides.slice_mut().swap(0, 1);
+            }
+        }
+        n => {
+            if let Some(min_stride_axis) = (0..n)
+                .filter(|&ax| dim[ax] > 1)
+                .min_by_key(|&ax| (strides[ax] as isize).abs())
+            {
+                let last = n - 1;
+                dim.slice_mut().swap(last, min_stride_axis);
+                strides.slice_mut().swap(last, min_stride_axis);
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::{

src/impl_methods.rs

@@ -18,8 +18,8 @@ use crate::arraytraits;
 use crate::dimension;
 use crate::dimension::IntoDimension;
 use crate::dimension::{
-    abs_index, axes_of, do_slice, merge_axes, offset_from_ptr_to_memory, size_of_shape_checked,
-    stride_offset, Axes,
+    abs_index, axes_of, do_slice, merge_axes, move_min_stride_axis_to_last,
+    offset_from_ptr_to_memory, size_of_shape_checked, stride_offset, Axes,
 };
 use crate::error::{self, ErrorKind, ShapeError};
 use crate::math_cell::MathCell;
@@ -1456,20 +1456,29 @@ where
     /// Return the array’s data as a slice if it is contiguous,
     /// return `None` otherwise.
     pub fn as_slice_memory_order_mut(&mut self) -> Option<&mut [A]>
+    where
+        S: DataMut,
+    {
+        self.try_as_slice_memory_order_mut().ok()
+    }
+
+    /// Return the array’s data as a slice if it is contiguous, otherwise
+    /// return `self` in the `Err` variant.
+    pub(crate) fn try_as_slice_memory_order_mut(&mut self) -> Result<&mut [A], &mut Self>
     where
         S: DataMut,
     {
         if self.is_contiguous() {
             self.ensure_unique();
             let offset = offset_from_ptr_to_memory(&self.dim, &self.strides);
             unsafe {
-                Some(slice::from_raw_parts_mut(
+                Ok(slice::from_raw_parts_mut(
                     self.ptr.offset(offset).as_ptr(),
                     self.len(),
                 ))
             }
         } else {
-            None
+            Err(self)
         }
     }
 
@@ -2070,27 +2079,7 @@ where
             slc.iter().fold(init, f)
         } else {
             let mut v = self.view();
-            // put the narrowest axis at the last position
-            match v.ndim() {
-                0 | 1 => {}
-                2 => {
-                    if self.len_of(Axis(1)) <= 1
-                        || self.len_of(Axis(0)) > 1
-                            && self.stride_of(Axis(0)).abs() < self.stride_of(Axis(1)).abs()
-                    {
-                        v.swap_axes(0, 1);
-                    }
-                }
-                n => {
-                    let last = n - 1;
-                    let narrow_axis = v
-                        .axes()
-                        .filter(|ax| ax.len() > 1)
-                        .min_by_key(|ax| ax.stride().abs())
-                        .map_or(last, |ax| ax.axis().index());
-                    v.swap_axes(last, narrow_axis);
-                }
-            }
+            move_min_stride_axis_to_last(&mut v.dim, &mut v.strides);
             v.into_elements_base().fold(init, f)
         }
     }
@@ -2200,12 +2189,20 @@ where
     /// Modify the array in place by calling `f` by mutable reference on each element.
     ///
     /// Elements are visited in arbitrary order.
-    pub fn map_inplace<F>(&mut self, f: F)
+    pub fn map_inplace<'a, F>(&'a mut self, f: F)
     where
         S: DataMut,
-        F: FnMut(&mut A),
-    {
-        self.unordered_foreach_mut(f);
+        A: 'a,
+        F: FnMut(&'a mut A),
+    {
+        match self.try_as_slice_memory_order_mut() {
+            Ok(slc) => slc.iter_mut().for_each(f),
+            Err(arr) => {
+                let mut v = arr.view_mut();
+                move_min_stride_axis_to_last(&mut v.dim, &mut v.strides);
+                v.into_elements_base().for_each(f);
+            }
+        }
     }
 
     /// Modify the array in place by calling `f` by **v**alue on each element.