Merge pull request #967 from rust-ndarray/shift-remove

Add method .remove_index(axis, index)

Author: bluss

src/impl_1d.rs

@@ -8,7 +8,10 @@
 
 //! Methods for one-dimensional arrays.
 use alloc::vec::Vec;
+use std::mem::MaybeUninit;
+
 use crate::imp_prelude::*;
+use crate::low_level_util::AbortIfPanic;
 
 /// # Methods For 1-D Arrays
 impl<A, S> ArrayBase<S, Ix1>
@@ -27,4 +30,35 @@ where
             crate::iterators::to_vec(self.iter().cloned())
         }
     }
+
+    /// Rotate the elements of the array by 1 element towards the front;
+    /// the former first element becomes the last.
+    pub(crate) fn rotate1_front(&mut self)
+    where
+        S: DataMut,
+    {
+        // use swapping to keep all elements initialized (as required by owned storage)
+        let mut lane_iter = self.iter_mut();
+        let mut dst = if let Some(dst) = lane_iter.next() { dst } else { return };
+
+        // Logically we do a circular swap here, all elements in a chain
+        // Using MaybeUninit to avoid unecessary writes in the safe swap solution
+        //
+        //  for elt in lane_iter {
+        //      std::mem::swap(dst, elt);
+        //      dst = elt;
+        //  }
+        //
+        let guard = AbortIfPanic(&"rotate1_front: temporarily moving out of owned value");
+        let mut slot = MaybeUninit::<A>::uninit();
+        unsafe {
+            slot.as_mut_ptr().copy_from_nonoverlapping(dst, 1);
+            for elt in lane_iter {
+                (dst as *mut A).copy_from_nonoverlapping(elt, 1);
+                dst = elt;
+            }
+            (dst as *mut A).copy_from_nonoverlapping(slot.as_ptr(), 1);
+        }
+        guard.defuse();
+    }
 }

src/impl_methods.rs

@@ -2444,6 +2444,30 @@ where
         }
     }
 
+    /// Remove the `index`th elements along `axis` and shift down elements from higher indexes.
+    ///
+    /// Note that this "removes" the elements by swapping them around to the end of the axis and
+    /// shortening the length of the axis; the elements are not deinitialized or dropped by this,
+    /// just moved out of view (this only matters for elements with ownership semantics). It's
+    /// similar to slicing an owned array in place.
+    ///
+    /// Decreases the length of `axis` by one.
+    ///
+    /// ***Panics*** if `axis` is out of bounds<br>
+    /// ***Panics*** if not `index < self.len_of(axis)`.
+    pub fn remove_index(&mut self, axis: Axis, index: usize)
+    where
+        S: DataOwned + DataMut,
+    {
+        assert!(index < self.len_of(axis), "index {} must be less than length of Axis({})",
+                index, axis.index());
+        let (_, mut tail) = self.view_mut().split_at(axis, index);
+        // shift elements to the front
+        Zip::from(tail.lanes_mut(axis)).for_each(|mut lane| lane.rotate1_front());
+        // then slice the axis in place to cut out the removed final element
+        self.slice_axis_inplace(axis, Slice::new(0, Some(-1), 1));
+    }
+
     /// Iterates over pairs of consecutive elements along the axis.
     ///
     /// The first argument to the closure is an element, and the second

src/lib.rs

@@ -205,6 +205,7 @@ mod shape_builder;
 mod slice;
 mod split_at;
 mod stacking;
+mod low_level_util;
 #[macro_use]
 mod zip;
 

src/low_level_util.rs

@@ -0,0 +1,40 @@
+// Copyright 2021 bluss and ndarray developers.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+
+/// Guard value that will abort if it is dropped.
+/// To defuse, this value must be forgotten before the end of the scope.
+///
+/// The string value is added to the message printed if aborting.
+#[must_use]
+pub(crate) struct AbortIfPanic(pub(crate) &'static &'static str);
+
+impl AbortIfPanic {
+    /// Defuse the AbortIfPanic guard. This *must* be done when finished.
+    #[inline]
+    pub(crate) fn defuse(self) {
+        std::mem::forget(self);
+    }
+}
+
+impl Drop for AbortIfPanic {
+    // The compiler should be able to remove this, if it can see through that there
+    // is no panic in the code section.
+    fn drop(&mut self) {
+        #[cfg(feature="std")]
+        {
+            eprintln!("ndarray: panic in no-panic section, aborting: {}", self.0);
+            std::process::abort()
+        }
+        #[cfg(not(feature="std"))]
+        {
+            // no-std uses panic-in-panic (should abort)
+            panic!("ndarray: panic in no-panic section, bailing out: {}", self.0);
+        }
+    }
+}

tests/array.rs

@@ -2397,3 +2397,79 @@ mod array_cow_tests {
         });
     }
 }
+
+#[test]
+fn test_remove_index() {
+    let mut a = arr2(&[[1, 2, 3],
+                       [4, 5, 6],
+                       [7, 8, 9],
+                       [10,11,12]]);
+    a.remove_index(Axis(0), 1);
+    a.remove_index(Axis(1), 2);
+    assert_eq!(a.shape(), &[3, 2]);
+    assert_eq!(a,
+        array![[1, 2],
+               [7, 8],
+               [10,11]]);
+
+    let mut a = arr2(&[[1, 2, 3],
+                       [4, 5, 6],
+                       [7, 8, 9],
+                       [10,11,12]]);
+    a.invert_axis(Axis(0));
+    a.remove_index(Axis(0), 1);
+    a.remove_index(Axis(1), 2);
+    assert_eq!(a.shape(), &[3, 2]);
+    assert_eq!(a,
+        array![[10,11],
+               [4, 5],
+               [1, 2]]);
+
+    a.remove_index(Axis(1), 1);
+
+    assert_eq!(a.shape(), &[3, 1]);
+    assert_eq!(a,
+        array![[10],
+               [4],
+               [1]]);
+    a.remove_index(Axis(1), 0);
+    assert_eq!(a.shape(), &[3, 0]);
+    assert_eq!(a,
+        array![[],
+               [],
+               []]);
+}
+
+#[should_panic(expected="must be less")]
+#[test]
+fn test_remove_index_oob1() {
+    let mut a = arr2(&[[1, 2, 3],
+                       [4, 5, 6],
+                       [7, 8, 9],
+                       [10,11,12]]);
+    a.remove_index(Axis(0), 4);
+}
+
+#[should_panic(expected="must be less")]
+#[test]
+fn test_remove_index_oob2() {
+    let mut a = array![[10], [4], [1]];
+    a.remove_index(Axis(1), 0);
+    assert_eq!(a.shape(), &[3, 0]);
+    assert_eq!(a,
+        array![[],
+               [],
+               []]);
+    a.remove_index(Axis(0), 1); // ok
+    assert_eq!(a,
+        array![[],
+               []]);
+    a.remove_index(Axis(1), 0); // oob
+}
+
+#[should_panic(expected="index out of bounds")]
+#[test]
+fn test_remove_index_oob3() {
+    let mut a = array![[10], [4], [1]];
+    a.remove_index(Axis(2), 0);
+}