Add support for individual indices when slicing

Author: jturner314

src/dimension/dimension_trait.rs

@@ -13,7 +13,7 @@ use std::ops::{Add, Sub, Mul, AddAssign, SubAssign, MulAssign};
 
 use itertools::{enumerate, zip};
 
-use {Ix, Ixs, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, Dim, Si, IxDynImpl};
+use {Ix, Ixs, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, Dim, SliceOrIndex, IxDynImpl};
 use IntoDimension;
 use RemoveAxis;
 use {ArrayView1, ArrayViewMut1};
@@ -52,14 +52,14 @@ pub trait Dimension : Clone + Eq + Debug + Send + Sync + Default +
     /// size, which you pass by reference. For the dynamic dimension it is
     /// a slice.
     ///
-    /// - For `Ix1`: `[Si; 1]`
-    /// - For `Ix2`: `[Si; 2]`
+    /// - For `Ix1`: `[SliceOrIndex; 1]`
+    /// - For `Ix2`: `[SliceOrIndex; 2]`
     /// - and so on..
-    /// - For `IxDyn`: `[Si]`
+    /// - For `IxDyn`: `[SliceOrIndex]`
     ///
-    /// The easiest way to create a `&SliceArg` is using the macro
-    /// [`s![]`](macro.s!.html).
-    type SliceArg: ?Sized + AsRef<[Si]>;
+    /// The easiest way to create a `&SliceInfo<SliceArg, Do>` is using the
+    /// [`s![]`](macro.s!.html) macro.
+    type SliceArg: ?Sized + AsRef<[SliceOrIndex]>;
     /// Pattern matching friendly form of the dimension value.
     ///
     /// - For `Ix1`: `usize`,
@@ -364,7 +364,7 @@ macro_rules! impl_insert_axis_array(
 
 impl Dimension for Dim<[Ix; 0]> {
     const NDIM: Option<usize> = Some(0);
-    type SliceArg = [Si; 0];
+    type SliceArg = [SliceOrIndex; 0];
     type Pattern = ();
     type Smaller = Self;
     type Larger = Ix1;
@@ -401,7 +401,7 @@ impl Dimension for Dim<[Ix; 0]> {
 
 impl Dimension for Dim<[Ix; 1]> {
     const NDIM: Option<usize> = Some(1);
-    type SliceArg = [Si; 1];
+    type SliceArg = [SliceOrIndex; 1];
     type Pattern = Ix;
     type Smaller = Ix0;
     type Larger = Ix2;
@@ -495,7 +495,7 @@ impl Dimension for Dim<[Ix; 1]> {
 
 impl Dimension for Dim<[Ix; 2]> {
     const NDIM: Option<usize> = Some(2);
-    type SliceArg = [Si; 2];
+    type SliceArg = [SliceOrIndex; 2];
     type Pattern = (Ix, Ix);
     type Smaller = Ix1;
     type Larger = Ix3;
@@ -631,7 +631,7 @@ impl Dimension for Dim<[Ix; 2]> {
 
 impl Dimension for Dim<[Ix; 3]> {
     const NDIM: Option<usize> = Some(3);
-    type SliceArg = [Si; 3];
+    type SliceArg = [SliceOrIndex; 3];
     type Pattern = (Ix, Ix, Ix);
     type Smaller = Ix2;
     type Larger = Ix4;
@@ -749,7 +749,7 @@ macro_rules! large_dim {
     ($n:expr, $name:ident, $pattern:ty, $larger:ty, { $($insert_axis:tt)* }) => (
         impl Dimension for Dim<[Ix; $n]> {
             const NDIM: Option<usize> = Some($n);
-            type SliceArg = [Si; $n];
+            type SliceArg = [SliceOrIndex; $n];
             type Pattern = $pattern;
             type Smaller = Dim<[Ix; $n - 1]>;
             type Larger = $larger;
@@ -801,7 +801,7 @@ large_dim!(6, Ix6, (Ix, Ix, Ix, Ix, Ix, Ix), IxDyn, {
 impl Dimension for IxDyn
 {
     const NDIM: Option<usize> = None;
-    type SliceArg = [Si];
+    type SliceArg = [SliceOrIndex];
     type Pattern = Self;
     type Smaller = Self;
     type Larger = Self;

src/dimension/mod.rs

@@ -214,7 +214,7 @@ pub fn do_sub<A, D: Dimension>(dims: &mut D, ptr: &mut *mut A, strides: &D,
 
 /// Compute the equivalent unsigned index given the axis length and signed index.
 #[inline]
-fn abs_index(len: Ix, index: Ixs) -> Ix {
+pub fn abs_index(len: Ix, index: Ixs) -> Ix {
     if index < 0 {
         len - (-index as Ix)
     } else {

src/impl_methods.rs

@@ -19,7 +19,7 @@ use dimension;
 use iterators;
 use error::{self, ShapeError, ErrorKind};
 use dimension::IntoDimension;
-use dimension::{axes_of, Axes, do_slice, merge_axes, stride_offset};
+use dimension::{abs_index, axes_of, Axes, do_slice, merge_axes, stride_offset};
 use iterators::{
     new_lanes,
     new_lanes_mut,
@@ -31,7 +31,8 @@ use zip::Zip;
 
 use {
     NdIndex,
-    Si,
+    SliceInfo,
+    SliceOrIndex
 };
 use iter::{
     AxisChunksIter,
@@ -207,53 +208,111 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
     }
 
 
-    /// Return a sliced array.
+    /// Return a sliced view of the array.
     ///
     /// See [*Slicing*](#slicing) for full documentation.
-    /// See also [`D::SliceArg`].
+    /// See also [`SliceInfo`] and [`D::SliceArg`].
     ///
+    /// [`SliceInfo`]: struct.SliceInfo.html
     /// [`D::SliceArg`]: trait.Dimension.html#associatedtype.SliceArg
     ///
-    /// **Panics** if an index is out of bounds or stride is zero.<br>
-    /// (**Panics** if `D` is `IxDyn` and `indexes` does not match the number of array axes.)
-    pub fn slice(&self, indexes: &D::SliceArg) -> ArrayView<A, D> {
-        let mut arr = self.view();
-        arr.slice_inplace(indexes);
-        arr
+    /// **Panics** if an index is out of bounds or step size is zero.<br>
+    /// (**Panics** if `D` is `IxDyn` and `info` does not match the number of array axes.)
+    pub fn slice<Do>(&self, info: &SliceInfo<D::SliceArg, Do>) -> ArrayView<A, Do>
+    where
+        Do: Dimension,
+    {
+        self.view().slice_move(info)
     }
 
     /// Return a sliced read-write view of the array.
     ///
-    /// See also [`D::SliceArg`].
+    /// See [*Slicing*](#slicing) for full documentation.
+    /// See also [`SliceInfo`] and [`D::SliceArg`].
     ///
+    /// [`SliceInfo`]: struct.SliceInfo.html
     /// [`D::SliceArg`]: trait.Dimension.html#associatedtype.SliceArg
     ///
-    /// **Panics** if an index is out of bounds or stride is zero.<br>
-    /// (**Panics** if `D` is `IxDyn` and `indexes` does not match the number of array axes.)
-    pub fn slice_mut(&mut self, indexes: &D::SliceArg) -> ArrayViewMut<A, D>
-        where S: DataMut
+    /// **Panics** if an index is out of bounds or step size is zero.<br>
+    /// (**Panics** if `D` is `IxDyn` and `info` does not match the number of array axes.)
+    pub fn slice_mut<Do>(&mut self, info: &SliceInfo<D::SliceArg, Do>) -> ArrayViewMut<A, Do>
+    where
+        Do: Dimension,
+        S: DataMut,
+    {
+        self.view_mut().slice_move(info)
+    }
+
+    /// Slice the array, possibly changing the number of dimensions.
+    ///
+    /// See [*Slicing*](#slicing) for full documentation.
+    /// See also [`SliceInfo`] and [`D::SliceArg`].
+    ///
+    /// [`SliceInfo`]: struct.SliceInfo.html
+    /// [`D::SliceArg`]: trait.Dimension.html#associatedtype.SliceArg
+    ///
+    /// **Panics** if an index is out of bounds or step size is zero.<br>
+    /// (**Panics** if `D` is `IxDyn` and `info` does not match the number of array axes.)
+    pub fn slice_move<Do>(mut self, info: &SliceInfo<D::SliceArg, Do>) -> ArrayBase<S, Do>
+    where
+        Do: Dimension,
     {
-        let mut arr = self.view_mut();
-        arr.slice_inplace(indexes);
-        arr
+        // Slice and subview in-place without changing the number of dimensions.
+        self.slice_inplace(&*info);
+
+        let indices: &[SliceOrIndex] = (**info).as_ref();
+
+        // Copy the dim and strides that remain after removing the subview axes.
+        let out_ndim = info.out_ndim();
+        let mut new_dim = Do::zero_index_with_ndim(out_ndim);
+        let mut new_strides = Do::zero_index_with_ndim(out_ndim);
+        izip!(self.dim.slice(), self.strides.slice(), indices)
+            .filter_map(|(d, s, slice_or_index)| match slice_or_index {
+                &SliceOrIndex::Slice(..) => Some((d, s)),
+                &SliceOrIndex::Index(_) => None,
+            })
+            .zip(izip!(new_dim.slice_mut(), new_strides.slice_mut()))
+            .for_each(|((d, s), (new_d, new_s))| {
+                *new_d = *d;
+                *new_s = *s;
+            });
+
+        ArrayBase {
+            ptr: self.ptr,
+            data: self.data,
+            dim: new_dim,
+            strides: new_strides,
+        }
     }
 
-    /// Slice the array’s view in place.
+    /// Slice the array in place without changing the number of dimensions.
+    ///
+    /// Note that [`&SliceInfo`](struct.SliceInfo.html) (produced by the
+    /// [`s![]`](macro.s!.html) macro) will usually coerce into `&D::SliceArg`
+    /// automatically, but in some cases (e.g. if `D` is `IxDyn`), you may need
+    /// to call `.as_ref()`.
     ///
+    /// See [*Slicing*](#slicing) for full documentation.
     /// See also [`D::SliceArg`].
     ///
     /// [`D::SliceArg`]: trait.Dimension.html#associatedtype.SliceArg
     ///
-    /// **Panics** if an index is out of bounds or stride is zero.<br>
-    /// (**Panics** if `D` is `IxDyn` and `indexes` does not match the number of array axes.)
-    pub fn slice_inplace(&mut self, indexes: &D::SliceArg) {
-        let indexes: &[Si] = indexes.as_ref();
-        assert_eq!(indexes.len(), self.ndim());
-        indexes
+    /// **Panics** if an index is out of bounds or step size is zero.<br>
+    /// (**Panics** if `D` is `IxDyn` and `indices` does not match the number of array axes.)
+    pub fn slice_inplace(&mut self, indices: &D::SliceArg) {
+        let indices: &[SliceOrIndex] = indices.as_ref();
+        assert_eq!(indices.len(), self.ndim());
+        indices
             .iter()
             .enumerate()
-            .for_each(|(axis, &Si(start, end, step))| {
-                self.slice_axis_inplace(Axis(axis), start, end, step)
+            .for_each(|(axis, slice_or_index)| match slice_or_index {
+                &SliceOrIndex::Slice(start, end, step) => {
+                    self.slice_axis_inplace(Axis(axis), start, end, step)
+                }
+                &SliceOrIndex::Index(index) => {
+                    let i_usize = abs_index(self.len_of(Axis(axis)), index);
+                    self.subview_inplace(Axis(axis), i_usize)
+                }
             });
     }
 

src/lib.rs

@@ -105,7 +105,7 @@ pub use dimension::NdIndex;
 pub use dimension::IxDynImpl;
 pub use indexes::{indices, indices_of};
 pub use error::{ShapeError, ErrorKind};
-pub use slice::{Si, S};
+pub use slice::{SliceInfo, SliceNextDim, SliceOrIndex};
 
 use iterators::Baseiter;
 use iterators::{ElementsBase, ElementsBaseMut, Iter, IterMut};
@@ -422,24 +422,30 @@ pub type Ixs = isize;
 ///
 /// You can use slicing to create a view of a subset of the data in
 /// the array. Slicing methods include [`.slice()`], [`.slice_mut()`],
-/// and [`.slice_inplace()`].
+/// [`.slice_move()`], and [`.slice_inplace()`].
 ///
 /// The slicing argument can be passed using the macro [`s![]`](macro.s!.html),
-/// which will be used in all examples. (The explicit form is a reference
-/// to a fixed size array of [`Si`]; see its docs for more information.)
+/// which will be used in all examples. (The explicit form is an instance of
+/// [`&SliceInfo`]; see its docs for more information.)
 ///
-/// [`Si`]: struct.Si.html
+/// [`&SliceInfo`]: struct.SliceInfo.html
+///
+/// If a range is used, the axis is preserved. If an index is used, a subview
+/// is taken with respect to the axis. See [*Subviews*](#subviews) for more
+/// information about subviews. Note that [`.slice_inplace()`] behaves like
+/// [`.subview_inplace()`] by preserving the number of dimensions.
 ///
 /// [`.slice()`]: #method.slice
 /// [`.slice_mut()`]: #method.slice_mut
+/// [`.slice_move()`]: #method.slice_move
 /// [`.slice_inplace()`]: #method.slice_inplace
 ///
 /// ```
 /// // import the s![] macro
 /// #[macro_use(s)]
 /// extern crate ndarray;
 ///
-/// use ndarray::arr3;
+/// use ndarray::{arr2, arr3};
 ///
 /// fn main() {
 ///
@@ -460,8 +466,6 @@ pub type Ixs = isize;
 /// // - Every element in each row: `..`
 ///
 /// let b = a.slice(s![.., 0..1, ..]);
-/// // without the macro, the explicit argument is `&[S, Si(0, Some(1), 1), S]`
-///
 /// let c = arr3(&[[[ 1,  2,  3]],
 ///                [[ 7,  8,  9]]]);
 /// assert_eq!(b, c);
@@ -476,14 +480,28 @@ pub type Ixs = isize;
 /// let e = arr3(&[[[ 6,  5,  4]],
 ///                [[12, 11, 10]]]);
 /// assert_eq!(d, e);
+/// assert_eq!(d.shape(), &[2, 1, 3]);
+///
+/// // Let’s create a slice while taking a subview with
+/// //
+/// // - Both submatrices of the greatest dimension: `..`
+/// // - The last row in each submatrix, removing that axis: `-1`
+/// // - Row elements in reverse order: `..;-1`
+/// let f = a.slice(s![.., -1, ..;-1]);
+/// let g = arr2(&[[ 6,  5,  4],
+///                [12, 11, 10]]);
+/// assert_eq!(f, g);
+/// assert_eq!(f.shape(), &[2, 3]);
 /// }
 /// ```
 ///
 /// ## Subviews
 ///
 /// Subview methods allow you to restrict the array view while removing one
 /// axis from the array. Subview methods include [`.subview()`],
-/// [`.subview_mut()`], [`.into_subview()`], and [`.subview_inplace()`].
+/// [`.subview_mut()`], [`.into_subview()`], and [`.subview_inplace()`]. You
+/// can also take a subview by using a single index instead of a range when
+/// slicing.
 ///
 /// Subview takes two arguments: `axis` and `index`.
 ///
@@ -493,7 +511,11 @@ pub type Ixs = isize;
 /// [`.subview_inplace()`]: #method.subview_inplace
 ///
 /// ```
-/// use ndarray::{arr3, aview2, Axis};
+/// #[macro_use(s)] extern crate ndarray;
+///
+/// use ndarray::{arr3, aview1, aview2, Axis};
+///
+/// # fn main() {
 ///
 /// // 2 submatrices of 2 rows with 3 elements per row, means a shape of `[2, 2, 3]`.
 ///
@@ -523,6 +545,11 @@ pub type Ixs = isize;
 ///
 /// assert_eq!(sub_col, aview2(&[[ 1,  4],
 ///                              [ 7, 10]]));
+///
+/// // You can take multiple subviews at once (and slice at the same time)
+/// let double_sub = a.slice(s![1, .., 0]);
+/// assert_eq!(double_sub, aview1(&[7, 10]));
+/// # }
 /// ```
 ///
 /// [`.subview_inplace()`] modifies the view in the same way as [`.subview()`],