From da71b98d2fbebda25a8788cf6288dbcbf69de5d9 Mon Sep 17 00:00:00 2001
From: Jim Turner <github@turner.link>
Date: Fri, 5 Feb 2021 17:23:32 -0500
Subject: [PATCH] Add slice_each_axis_inplace method

---
 src/impl_methods.rs | 43 +++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs          | 11 +++++++++++
 2 files changed, 54 insertions(+)

diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index 9eea42d15..4ad63264b 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -25,6 +25,7 @@ use crate::error::{self, ErrorKind, ShapeError};
 use crate::math_cell::MathCell;
 use crate::itertools::zip;
 use crate::zip::Zip;
+use crate::AxisDescription;
 
 use crate::iter::{
     AxisChunksIter, AxisChunksIterMut, AxisIter, AxisIterMut, ExactChunks, ExactChunksMut,
@@ -511,6 +512,48 @@ where
         debug_assert!(self.pointer_is_inbounds());
     }
 
+    /// Slice the array in place, with a closure specifying the slice for each
+    /// axis.
+    ///
+    /// This is especially useful for code which is generic over the
+    /// dimensionality of the array.
+    ///
+    /// **Panics** if an index is out of bounds or step size is zero.
+    ///
+    /// ```
+    /// use ndarray::{s, Array2, ArrayBase, ArrayView, Data, Dimension, Slice};
+    ///
+    /// fn view_halved_axes<S, D>(arr: &ArrayBase<S, D>) -> ArrayView<'_, S::Elem, D>
+    /// where
+    ///     S: Data,
+    ///     D: Dimension,
+    /// {
+    ///
+    ///     let mut view = arr.view();
+    ///     view.slice_each_axis_inplace(|ax| Slice::from(0..ax.len() / 2));
+    ///     view
+    /// }
+    ///
+    /// let a = Array2::<f32>::eye(8);
+    /// let v = view_halved_axes(&a);
+    /// assert_eq!(v, a.slice(s![..4, ..4]));
+    /// ```
+    pub fn slice_each_axis_inplace<F>(&mut self, mut f: F)
+    where
+        F: FnMut(AxisDescription) -> Slice,
+    {
+        (0..self.ndim()).for_each(|ax| {
+            self.slice_axis_inplace(
+                Axis(ax),
+                f(AxisDescription(
+                    Axis(ax),
+                    self.dim[ax],
+                    self.strides[ax] as isize,
+                )),
+            )
+        })
+    }
+
     /// Return a reference to the element at `index`, or return `None`
     /// if the index is out of bounds.
     ///
diff --git a/src/lib.rs b/src/lib.rs
index 0c7caf735..4f2936c5b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -503,6 +503,17 @@ pub type Ixs = isize;
 /// [`.slice_move()`]: #method.slice_move
 /// [`.slice_collapse()`]: #method.slice_collapse
 ///
+/// When slicing arrays with generic dimensionality, creating an instance of
+/// [`&SliceInfo`] to pass to the multi-axis slicing methods like [`.slice()`]
+/// is awkward. In these cases, it's usually more convenient to create a view
+/// and then slice individual axes of the view using methods such as
+/// [`.slice_axis_inplace()`], [`.collapse_axis()`], and
+/// [`.slice_each_axis_inplace()`].
+///
+/// [`.slice_axis_inplace()`]: #method.slice_axis_inplace
+/// [`.collapse_axis()`]: #method.collapse_axis
+/// [`.slice_each_axis_inplace()`]: #method.slice_each_axis_inplace
+///
 /// It's possible to take multiple simultaneous *mutable* slices with
 /// [`.multi_slice_mut()`] or (for [`ArrayViewMut`] only)
 /// [`.multi_slice_move()`].