Make slice_collapse return Err(_) for NewAxis

Author: jturner314

examples/axis_ops.rs

@@ -51,7 +51,7 @@ fn main() {
     }
     a.swap_axes(0, 1);
     a.swap_axes(0, 2);
-    a.slice_collapse(s![.., ..;-1, ..]);
+    a.slice_collapse(s![.., ..;-1, ..]).unwrap();
     regularize(&mut a).ok();
 
     let mut b = Array::<u8, _>::zeros((2, 3, 4));
@@ -68,6 +68,6 @@ fn main() {
     for (i, elt) in (0..).zip(&mut a) {
         *elt = i;
     }
-    a.slice_collapse(s![..;-1, ..;2, ..]);
+    a.slice_collapse(s![..;-1, ..;2, ..]).unwrap();
     regularize(&mut a).ok();
 }

serialization-tests/tests/serialize.rs

@@ -46,7 +46,7 @@ fn serial_many_dim_serde() {
     {
         // Test a sliced array.
         let mut a = ArcArray::linspace(0., 31., 32).reshape((2, 2, 2, 4));
-        a.slice_collapse(s![..;-1, .., .., ..2]);
+        a.slice_collapse(s![..;-1, .., .., ..2]).unwrap();
         let serial = serde_json::to_string(&a).unwrap();
         println!("Encode {:?} => {:?}", a, serial);
         let res = serde_json::from_str::<ArcArray<f32, _>>(&serial);
@@ -156,7 +156,7 @@ fn serial_many_dim_serde_msgpack() {
     {
         // Test a sliced array.
         let mut a = ArcArray::linspace(0., 31., 32).reshape((2, 2, 2, 4));
-        a.slice_collapse(s![..;-1, .., .., ..2]);
+        a.slice_collapse(s![..;-1, .., .., ..2]).unwrap();
 
         let mut buf = Vec::new();
         serde::Serialize::serialize(&a, &mut rmp_serde::Serializer::new(&mut buf))
@@ -209,7 +209,7 @@ fn serial_many_dim_ron() {
     {
         // Test a sliced array.
         let mut a = ArcArray::linspace(0., 31., 32).reshape((2, 2, 2, 4));
-        a.slice_collapse(s![..;-1, .., .., ..2]);
+        a.slice_collapse(s![..;-1, .., .., ..2]).unwrap();
 
         let a_s = ron_serialize(&a).unwrap();
 

src/impl_methods.rs

@@ -461,13 +461,15 @@ where
 
     /// Slice the array in place without changing the number of dimensions.
     ///
-    /// Note that `NewAxis` elements in `info` are ignored.
+    /// If there are any `NewAxis` elements in `info`, slicing is performed
+    /// using the other elements in `info` (i.e. ignoring the `NewAxis`
+    /// elements), and `Err(_)` is returned to notify the caller.
     ///
     /// See [*Slicing*](#slicing) for full documentation.
     ///
     /// **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_collapse<I>(&mut self, info: &I)
+    pub fn slice_collapse<I>(&mut self, info: &I) -> Result<(), ShapeError>
     where
         I: CanSlice<D> + ?Sized,
     {
@@ -476,20 +478,28 @@ where
             self.ndim(),
             "The input dimension of `info` must match the array to be sliced.",
         );
+        let mut new_axis_in_info = false;
         let mut axis = 0;
         info.as_ref().iter().for_each(|&ax_info| match ax_info {
-                AxisSliceInfo::Slice { start, end, step } => {
-                    self.slice_axis_inplace(Axis(axis), Slice { start, end, step });
-                    axis += 1;
-                }
-                AxisSliceInfo::Index(index) => {
-                    let i_usize = abs_index(self.len_of(Axis(axis)), index);
-                    self.collapse_axis(Axis(axis), i_usize);
-                    axis += 1;
-                }
-                AxisSliceInfo::NewAxis => {}
-            });
+            AxisSliceInfo::Slice { start, end, step } => {
+                self.slice_axis_inplace(Axis(axis), Slice { start, end, step });
+                axis += 1;
+            }
+            AxisSliceInfo::Index(index) => {
+                let i_usize = abs_index(self.len_of(Axis(axis)), index);
+                self.collapse_axis(Axis(axis), i_usize);
+                axis += 1;
+            }
+            AxisSliceInfo::NewAxis => {
+                new_axis_in_info = true;
+            }
+        });
         debug_assert_eq!(axis, self.ndim());
+        if new_axis_in_info {
+            Err(ShapeError::from_kind(ErrorKind::IncompatibleShape))
+        } else {
+            Ok(())
+        }
     }
 
     /// Return a view of the array, sliced along the specified axis.

tests/array.rs

@@ -103,10 +103,10 @@ fn test_slice_ix0() {
 #[test]
 fn test_slice_edge_cases() {
     let mut arr = Array3::<u8>::zeros((3, 4, 5));
-    arr.slice_collapse(s![0..0;-1, .., ..]);
+    arr.slice_collapse(s![0..0;-1, .., ..]).unwrap();
     assert_eq!(arr.shape(), &[0, 4, 5]);
     let mut arr = Array2::<u8>::from_shape_vec((1, 1).strides((10, 1)), vec![5]).unwrap();
-    arr.slice_collapse(s![1..1, ..]);
+    arr.slice_collapse(s![1..1, ..]).unwrap();
     assert_eq!(arr.shape(), &[0, 1]);
 }
 
@@ -201,7 +201,7 @@ fn test_slice_array_fixed() {
     arr.slice(info);
     arr.slice_mut(info);
     arr.view().slice_move(info);
-    arr.view().slice_collapse(info);
+    arr.view().slice_collapse(info).unwrap_err();
 }
 
 #[test]
@@ -211,7 +211,7 @@ fn test_slice_dyninput_array_fixed() {
     arr.slice(info);
     arr.slice_mut(info);
     arr.view().slice_move(info);
-    arr.view().slice_collapse(info);
+    arr.view().slice_collapse(info).unwrap_err();
 }
 
 #[test]
@@ -227,7 +227,7 @@ fn test_slice_array_dyn() {
     arr.slice(info);
     arr.slice_mut(info);
     arr.view().slice_move(info);
-    arr.view().slice_collapse(info);
+    arr.view().slice_collapse(info).unwrap_err();
 }
 
 #[test]
@@ -243,7 +243,7 @@ fn test_slice_dyninput_array_dyn() {
     arr.slice(info);
     arr.slice_mut(info);
     arr.view().slice_move(info);
-    arr.view().slice_collapse(info);
+    arr.view().slice_collapse(info).unwrap_err();
 }
 
 #[test]
@@ -259,7 +259,7 @@ fn test_slice_dyninput_vec_fixed() {
     arr.slice(info);
     arr.slice_mut(info);
     arr.view().slice_move(info);
-    arr.view().slice_collapse(info);
+    arr.view().slice_collapse(info).unwrap_err();
 }
 
 #[test]
@@ -275,7 +275,7 @@ fn test_slice_dyninput_vec_dyn() {
     arr.slice(info);
     arr.slice_mut(info);
     arr.view().slice_move(info);
-    arr.view().slice_collapse(info);
+    arr.view().slice_collapse(info).unwrap_err();
 }
 
 #[test]
@@ -324,31 +324,31 @@ fn test_slice_collapse_with_indices() {
 
     {
         let mut vi = arr.view();
-        vi.slice_collapse(s![NewAxis, 1.., 2, ..;2]);
+        vi.slice_collapse(s![NewAxis, 1.., 2, ..;2]).unwrap_err();
         assert_eq!(vi.shape(), &[2, 1, 2]);
         assert!(vi
             .iter()
             .zip(arr.slice(s![1.., 2..3, ..;2]).iter())
             .all(|(a, b)| a == b));
 
         let mut vi = arr.view();
-        vi.slice_collapse(s![1, NewAxis, 2, ..;2]);
+        vi.slice_collapse(s![1, NewAxis, 2, ..;2]).unwrap_err();
         assert_eq!(vi.shape(), &[1, 1, 2]);
         assert!(vi
             .iter()
             .zip(arr.slice(s![1..2, 2..3, ..;2]).iter())
             .all(|(a, b)| a == b));
 
         let mut vi = arr.view();
-        vi.slice_collapse(s![1, 2, NewAxis, 3]);
+        vi.slice_collapse(s![1, 2, 3]).unwrap();
         assert_eq!(vi.shape(), &[1, 1, 1]);
         assert_eq!(vi, Array3::from_elem((1, 1, 1), arr[(1, 2, 3)]));
     }
 
     // Do it to the ArcArray itself
     let elem = arr[(1, 2, 3)];
     let mut vi = arr;
-    vi.slice_collapse(s![1, 2, 3, NewAxis]);
+    vi.slice_collapse(s![1, 2, 3, NewAxis]).unwrap_err();
     assert_eq!(vi.shape(), &[1, 1, 1]);
     assert_eq!(vi, Array3::from_elem((1, 1, 1), elem));
 }
@@ -567,7 +567,7 @@ fn test_cow() {
     assert_eq!(n[[0, 1]], 0);
     assert_eq!(n.get((0, 1)), Some(&0));
     let mut rev = mat.reshape(4);
-    rev.slice_collapse(s![..;-1]);
+    rev.slice_collapse(s![..;-1]).unwrap();
     assert_eq!(rev[0], 4);
     assert_eq!(rev[1], 3);
     assert_eq!(rev[2], 2);
@@ -591,7 +591,7 @@ fn test_cow_shrink() {
     // mutation shrinks the array and gives it different strides
     //
     let mut mat = ArcArray::zeros((2, 3));
-    //mat.slice_collapse(s![.., ..;2]);
+    //mat.slice_collapse(s![.., ..;2]).unwrap();
     mat[[0, 0]] = 1;
     let n = mat.clone();
     mat[[0, 1]] = 2;
@@ -606,7 +606,7 @@ fn test_cow_shrink() {
     assert_eq!(n.get((0, 1)), Some(&0));
     // small has non-C strides this way
     let mut small = mat.reshape(6);
-    small.slice_collapse(s![4..;-1]);
+    small.slice_collapse(s![4..;-1]).unwrap();
     assert_eq!(small[0], 6);
     assert_eq!(small[1], 5);
     let before = small.clone();
@@ -886,7 +886,7 @@ fn assign() {
     let mut a = arr2(&[[1, 2], [3, 4]]);
     {
         let mut v = a.view_mut();
-        v.slice_collapse(s![..1, ..]);
+        v.slice_collapse(s![..1, ..]).unwrap();
         v.fill(0);
     }
     assert_eq!(a, arr2(&[[0, 0], [3, 4]]));
@@ -1093,7 +1093,7 @@ fn owned_array_discontiguous_drop() {
             .collect();
         let mut a = Array::from_shape_vec((2, 6), v).unwrap();
         // discontiguous and non-zero offset
-        a.slice_collapse(s![.., 1..]);
+        a.slice_collapse(s![.., 1..]).unwrap();
     }
     // each item was dropped exactly once
     itertools::assert_equal(set.borrow().iter().cloned(), 0..12);
@@ -1792,7 +1792,7 @@ fn to_owned_memory_order() {
 #[test]
 fn to_owned_neg_stride() {
     let mut c = arr2(&[[1, 2, 3], [4, 5, 6]]);
-    c.slice_collapse(s![.., ..;-1]);
+    c.slice_collapse(s![.., ..;-1]).unwrap();
     let co = c.to_owned();
     assert_eq!(c, co);
     assert_eq!(c.strides(), co.strides());
@@ -1801,7 +1801,7 @@ fn to_owned_neg_stride() {
 #[test]
 fn discontiguous_owned_to_owned() {
     let mut c = arr2(&[[1, 2, 3], [4, 5, 6]]);
-    c.slice_collapse(s![.., ..;2]);
+    c.slice_collapse(s![.., ..;2]).unwrap();
 
     let co = c.to_owned();
     assert_eq!(c.strides(), &[3, 2]);
@@ -2062,10 +2062,10 @@ fn test_accumulate_axis_inplace_nonstandard_layout() {
 fn test_to_vec() {
     let mut a = arr2(&[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]);
 
-    a.slice_collapse(s![..;-1, ..]);
+    a.slice_collapse(s![..;-1, ..]).unwrap();
     assert_eq!(a.row(3).to_vec(), vec![1, 2, 3]);
     assert_eq!(a.column(2).to_vec(), vec![12, 9, 6, 3]);
-    a.slice_collapse(s![.., ..;-1]);
+    a.slice_collapse(s![.., ..;-1]).unwrap();
     assert_eq!(a.row(3).to_vec(), vec![3, 2, 1]);
 }
 
@@ -2081,7 +2081,7 @@ fn test_array_clone_unalias() {
 #[test]
 fn test_array_clone_same_view() {
     let mut a = Array::from_iter(0..9).into_shape((3, 3)).unwrap();
-    a.slice_collapse(s![..;-1, ..;-1]);
+    a.slice_collapse(s![..;-1, ..;-1]).unwrap();
     let b = a.clone();
     assert_eq!(a, b);
 }

tests/iterators.rs

@@ -332,7 +332,7 @@ fn axis_iter_zip_partially_consumed_discontiguous() {
     while iter.next().is_some() {
         consumed += 1;
         let mut b = Array::zeros((a.len() - consumed) * 2);
-        b.slice_collapse(s![..;2]);
+        b.slice_collapse(s![..;2]).unwrap();
         Zip::from(&mut b).and(iter.clone()).for_each(|b, a| *b = a[()]);
         assert_eq!(a.slice(s![consumed..]), b);
     }
@@ -519,7 +519,7 @@ fn axis_iter_mut_zip_partially_consumed_discontiguous() {
             iter.next();
         }
         let mut b = Array::zeros(remaining * 2);
-        b.slice_collapse(s![..;2]);
+        b.slice_collapse(s![..;2]).unwrap();
         Zip::from(&mut b).and(iter).for_each(|b, a| *b = a[()]);
         assert_eq!(a.slice(s![consumed..]), b);
     }