Implement ArrayVec::remove() and ::drain()

src/lib.rs

@@ -17,7 +17,9 @@ use std::hash::{Hash, Hasher};
 use std::fmt;
 
 mod array;
+mod misc;
 pub use array::Array;
+pub use misc::RangeArgument;
 
 
 unsafe fn new_array<A: Array>() -> A {
@@ -169,9 +171,7 @@ impl<A: Array> ArrayVec<A> {
     ///
     /// let mut array = ArrayVec::from([1, 2, 3]);
     ///
-    /// let elt = array.swap_remove(0);
+    /// assert_eq!(array.swap_remove(0), Some(1));
-    ///
-    /// assert_eq!(elt, Some(1));
     /// assert_eq!(&array[..], &[3, 2]);
     ///
     /// assert_eq!(array.swap_remove(10), None);
@@ -184,6 +184,80 @@ impl<A: Array> ArrayVec<A> {
         self.swap(index, len - 1);
         self.pop()
     }
+
+    /// Remove the element at **index** and shift the following elements down.
+    ///
+    /// Return **Some(** *element* **)** if the index is in bounds, else **None**.
+    ///
+    /// ## Examples
+    /// ```
+    /// use arrayvec::ArrayVec;
+    ///
+    /// let mut array = ArrayVec::from([1, 2, 3]);
+    ///
+    /// assert_eq!(array.remove(0), Some(1));
+    /// assert_eq!(&array[..], &[2, 3]);
+    ///
+    /// assert_eq!(array.remove(10), None);
+    /// ```
+    pub fn remove(&mut self, index: usize) -> Option<A::Item> {
+        if index >= self.len() {
+            None
+        } else {
+            self.drain(index..index + 1).next()
+        }
+    }
+
+    /// Create a draining iterator that removes the specified range in the vector
+    /// and yields the removed items from start to end. The element range is
+    /// removed even if the iterator is not consumed until the end.
+    ///
+    /// Note: It is unspecified how many elements are removed from the vector,
+    /// if the `Drain` value is leaked.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the starting point is greater than the end point or if
+    /// the end point is greater than the length of the vector.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use arrayvec::ArrayVec;
+    ///
+    /// let mut v = ArrayVec::from([1, 2, 3]);
+    /// let u: Vec<_> = v.drain(0..2).collect();
+    /// assert_eq!(&v[..], &[3]);
+    /// assert_eq!(&u[..], &[1, 2]);
+    /// ```
+    pub fn drain<R: RangeArgument>(&mut self, range: R) -> Drain<A> {
+        // Memory safety
+        //
+        // When the Drain is first created, it shortens the length of
+        // the source vector to make sure no uninitalized or moved-from elements
+        // are accessible at all if the Drain's destructor never gets to run.
+        //
+        // Drain will ptr::read out the values to remove.
+        // When finished, remaining tail of the vec is copied back to cover
+        // the hole, and the vector length is restored to the new length.
+        //
+        let len = self.len();
+        let start = range.start().unwrap_or(0);
+        let end = range.end().unwrap_or(len);
+        // bounds check happens here
+        let range_slice: *const _ = &self[start..end];
+
+        unsafe {
+            // set self.vec length's to start, to be safe in case Drain is leaked
+            self.len = start as u8;
+            Drain {
+                tail_start: end,
+                tail_len: len - end,
+                iter: (*range_slice).iter(),
+                vec: self as *mut _,
+            }
+        }
+    }
 }
 
 impl<A: Array> Deref for ArrayVec<A> {
@@ -283,7 +357,7 @@ impl<A: Array> IntoIterator for ArrayVec<A> {
 }
 
 
-/// By-value iterator for ArrayVec.
+/// By-value iterator for **ArrayVec**.
 pub struct IntoIter<A: Array> {
     index: u8,
     v: ArrayVec<A>,
@@ -338,6 +412,83 @@ impl<A: Array> Drop for IntoIter<A> {
     }
 }
 
+/// A draining iterator for **ArrayVec**.
+pub struct Drain<'a, A> 
+    where A: Array,
+          A::Item: 'a,
+{
+    /// Index of tail to preserve
+    tail_start: usize,
+    /// Length of tail
+    tail_len: usize,
+    /// Current remaining range to remove
+    iter: slice::Iter<'a, A::Item>,
+    vec: *mut ArrayVec<A>,
+}
+
+unsafe impl<'a, A: Array + Sync> Sync for Drain<'a, A> {}
+unsafe impl<'a, A: Array + Send> Send for Drain<'a, A> {}
+
+impl<'a, A: Array> Iterator for Drain<'a, A>
+    where A::Item: 'a,
+{
+    type Item = A::Item;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next().map(|elt|
+            unsafe {
+                ptr::read(elt as *const _)
+            }
+        )
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+impl<'a, A: Array> DoubleEndedIterator for Drain<'a, A>
+    where A::Item: 'a,
+{
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.iter.next_back().map(|elt|
+            unsafe {
+                ptr::read(elt as *const _)
+            }
+        )
+    }
+}
+
+impl<'a, A: Array> ExactSizeIterator for Drain<'a, A> where A::Item: 'a {}
+
+impl<'a, A: Array> Drop for Drain<'a, A> 
+    where A::Item: 'a
+{
+    fn drop(&mut self) {
+        // exhaust self first
+        while let Some(_) = self.next() { }
+
+        if self.tail_len > 0 {
+            unsafe {
+                let source_vec = &mut *self.vec;
+                // memmove back untouched tail, update to new length
+                let start = source_vec.len();
+                let tail = self.tail_start;
+                let src = source_vec.as_ptr().offset(tail as isize);
+                let dst = source_vec.as_mut_ptr().offset(start as isize);
+                ptr::copy(src, dst, self.tail_len);
+                source_vec.len = (start + self.tail_len) as u8;
+            }
+        }
+    }
+}
+
+
+
+
 /// Extend the **ArrayVec** with an iterator.
 /// 
 /// Does not extract more items than there is space for. No error
@@ -520,3 +671,28 @@ fn test_compact_size() {
     println!("{}", mem::size_of::<QuadArray>());
     assert!(mem::size_of::<QuadArray>() <= 24);
 }
+
+#[test]
+fn test_drain() {
+    let mut v = ArrayVec::from([0; 8]);
+    v.pop();
+    v.drain(0..7);
+    assert_eq!(&v[..], &[]);
+
+    v.extend(0..);
+    v.drain(1..4);
+    assert_eq!(&v[..], &[0, 4, 5, 6, 7]);
+    let u: ArrayVec<[_; 3]> = v.drain(1..4).rev().collect();
+    assert_eq!(&u[..], &[6, 5, 4]);
+    assert_eq!(&v[..], &[0, 7]);
+    v.drain(..);
+    assert_eq!(&v[..], &[]);
+}
+
+#[test]
+#[should_panic]
+fn test_drain_oob() {
+    let mut v = ArrayVec::from([0; 8]);
+    v.pop();
+    v.drain(0..8);
+}

src/misc.rs

@@ -0,0 +1,34 @@
+use std::ops::{
+    RangeFull,
+    RangeFrom,
+    RangeTo,
+    Range,
+};
+
+/// **RangeArgument** is implemented by Rust's built-in range types, produced
+/// by range syntax like `..`, `a..`, `..b` or `c..d`.
+pub trait RangeArgument {
+    #[doc(hidden)]
+    /// Start index (inclusive)
+    fn start(&self) -> Option<usize> { None }
+    #[doc(hidden)]
+    /// End index (exclusive)
+    fn end(&self) -> Option<usize> { None }
+}
+
+
+impl RangeArgument for RangeFull {}
+
+impl RangeArgument for RangeFrom<usize> {
+    fn start(&self) -> Option<usize> { Some(self.start) }
+}
+
+impl RangeArgument for RangeTo<usize> {
+    fn end(&self) -> Option<usize> { Some(self.end) }
+}
+
+impl RangeArgument for Range<usize> {
+    fn start(&self) -> Option<usize> { Some(self.start) }
+    fn end(&self) -> Option<usize> { Some(self.end) }
+}
+