collections: Implement vec::drain(range) according to RFC 574

Old `.drain()` on vec is performed using `.drain(..)` now.

[breaking-change]

Author: Ulrik Sverdrup

src/libcollections/binary_heap.rs

@@ -546,7 +546,7 @@ impl<T: Ord> BinaryHeap<T> {
     #[unstable(feature = "collections",
                reason = "matches collection reform specification, waiting for dust to settle")]
     pub fn drain(&mut self) -> Drain<T> {
-        Drain { iter: self.data.drain() }
+        Drain { iter: self.data.drain(..) }
     }
 
     /// Drops all items from the binary heap.

src/libcollections/vec.rs

@@ -69,6 +69,8 @@ use core::usize;
 
 use borrow::{Cow, IntoCow};
 
+use super::range::RangeArgument;
+
 // FIXME- fix places which assume the max vector allowed has memory usize::MAX.
 static MAX_MEMORY_SIZE: usize = isize::MAX as usize;
 
@@ -718,36 +720,61 @@ impl<T> Vec<T> {
         unsafe { other.set_len(0); }
     }
 
-    /// Creates a draining iterator that clears the `Vec` and iterates over
-    /// the removed items from start to end.
+    /// 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
     ///
     /// ```
     /// # #![feature(collections)]
-    /// let mut v = vec!["a".to_string(), "b".to_string()];
-    /// for s in v.drain() {
-    ///     // s has type String, not &String
-    ///     println!("{}", s);
-    /// }
-    /// assert!(v.is_empty());
+    ///
+    /// // Draining using `..` clears the whole vector.
+    /// let mut v = vec![1, 2, 3];
+    /// let u: Vec<_> = v.drain(..).collect();
+    /// assert_eq!(v, &[]);
+    /// assert_eq!(u, &[1, 2, 3]);
     /// ```
-    #[inline]
     #[unstable(feature = "collections",
-               reason = "matches collection reform specification, waiting for dust to settle")]
-    pub fn drain(&mut self) -> Drain<T> {
+               reason = "recently added, matches RFC")]
+    pub fn drain<R>(&mut self, range: R) -> Drain<T> where
+        R: RangeArgument<usize>,
+    {
+        // 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);
+        assert!(start <= end);
+        assert!(end <= len);
+
         unsafe {
-            let begin = *self.ptr as *const T;
-            let end = if mem::size_of::<T>() == 0 {
-                (*self.ptr as usize + self.len()) as *const T
-            } else {
-                (*self.ptr).offset(self.len() as isize) as *const T
-            };
-            self.set_len(0);
+            // set self.vec length's to start, to be safe in case Drain is leaked
+            self.set_len(start);
+            let range_slice = slice::from_raw_parts(self.as_ptr().offset(start as isize),
+                                                    end - start);
             Drain {
-                ptr: begin,
-                end: end,
-                marker: PhantomData,
+                tail_start: end,
+                tail_len: len - end,
+                iter: range_slice.iter(),
+                vec: self as *mut _,
+                _marker: PhantomData,
             }
         }
     }
@@ -1799,95 +1826,92 @@ impl<T> Drop for IntoIter<T> {
     }
 }
 
-/// An iterator that drains a vector.
-#[unsafe_no_drop_flag]
-#[unstable(feature = "collections",
-           reason = "recently added as part of collections reform 2")]
-pub struct Drain<'a, T:'a> {
-    ptr: *const T,
-    end: *const T,
-    marker: PhantomData<&'a T>,
+/// A draining iterator for `Vec<T>`.
+#[unstable(feature = "collections", reason = "recently added")]
+pub struct Drain<'a, T: 'a>
+{
+    /// Index of tail to preserve
+    tail_start: usize,
+    /// Length of tail
+    tail_len: usize,
+    /// Current remaining range to remove
+    iter: slice::Iter<'a, T>,
+    vec: *mut Vec<T>,
+    _marker: PhantomData<&'a mut Vec<T>>,
 }
 
 unsafe impl<'a, T: Sync> Sync for Drain<'a, T> {}
 unsafe impl<'a, T: Send> Send for Drain<'a, T> {}
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> Iterator for Drain<'a, T> {
+impl<'a, T> Iterator for Drain<'a, T>
+{
     type Item = T;
 
     #[inline]
-    fn next(&mut self) -> Option<T> {
-        unsafe {
-            if self.ptr == self.end {
-                None
-            } else {
-                if mem::size_of::<T>() == 0 {
-                    // purposefully don't use 'ptr.offset' because for
-                    // vectors with 0-size elements this would return the
-                    // same pointer.
-                    self.ptr = mem::transmute(self.ptr as usize + 1);
-
-                    // Use a non-null pointer value
-                    Some(ptr::read(EMPTY as *mut T))
-                } else {
-                    let old = self.ptr;
-                    self.ptr = self.ptr.offset(1);
-
-                    Some(ptr::read(old))
+    fn next(&mut self) -> Option<T>
+    {
+        match self.iter.next() {
+            None => None,
+            Some(elt) => {
+                unsafe {
+                    Some(ptr::read(elt as *const _))
                 }
             }
         }
     }
 
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        let diff = (self.end as usize) - (self.ptr as usize);
-        let size = mem::size_of::<T>();
-        let exact = diff / (if size == 0 {1} else {size});
-        (exact, Some(exact))
+    fn size_hint(&self) -> (usize, Option<usize>)
+    {
+        self.iter.size_hint()
     }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
+impl<'a, T> DoubleEndedIterator for Drain<'a, T>
+{
     #[inline]
-    fn next_back(&mut self) -> Option<T> {
-        unsafe {
-            if self.end == self.ptr {
-                None
-            } else {
-                if mem::size_of::<T>() == 0 {
-                    // See above for why 'ptr.offset' isn't used
-                    self.end = mem::transmute(self.end as usize - 1);
-
-                    // Use a non-null pointer value
-                    Some(ptr::read(EMPTY as *mut T))
-                } else {
-                    self.end = self.end.offset(-1);
-
-                    Some(ptr::read(self.end))
+    fn next_back(&mut self) -> Option<T>
+    {
+        match self.iter.next_back() {
+            None => None,
+            Some(elt) => {
+                unsafe {
+                    Some(ptr::read(elt as *const _))
                 }
             }
         }
     }
 }
 
-#[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> ExactSizeIterator for Drain<'a, T> {}
-
 #[unsafe_destructor]
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> Drop for Drain<'a, T> {
-    fn drop(&mut self) {
-        // self.ptr == self.end == mem::POST_DROP_USIZE if drop has already been called,
-        // so we can use #[unsafe_no_drop_flag].
-
-        // destroy the remaining elements
-        for _x in self.by_ref() {}
+impl<'a, T> Drop for Drain<'a, T>
+{
+    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.set_len(start + self.tail_len);
+            }
+        }
     }
 }
 
+
+#[stable(feature = "rust1", since = "1.0.0")]
+impl<'a, T> ExactSizeIterator for Drain<'a, T> {}
+
 ////////////////////////////////////////////////////////////////////////////////
 // Conversion from &[T] to &Vec<T>
 ////////////////////////////////////////////////////////////////////////////////

src/libcollections/vec_map.rs

@@ -418,7 +418,7 @@ impl<V> VecMap<V> {
         }
         let filter: fn((usize, Option<V>)) -> Option<(usize, V)> = filter; // coerce to fn ptr
 
-        Drain { iter: self.v.drain().enumerate().filter_map(filter) }
+        Drain { iter: self.v.drain(..).enumerate().filter_map(filter) }
     }
 
     /// Returns the number of elements in the map.

src/libcollectionstest/vec.rs

@@ -460,7 +460,7 @@ fn test_move_items_zero_sized() {
 fn test_drain_items() {
     let mut vec = vec![1, 2, 3];
     let mut vec2 = vec![];
-    for i in vec.drain() {
+    for i in vec.drain(..) {
         vec2.push(i);
     }
     assert_eq!(vec, []);
@@ -471,7 +471,7 @@ fn test_drain_items() {
 fn test_drain_items_reverse() {
     let mut vec = vec![1, 2, 3];
     let mut vec2 = vec![];
-    for i in vec.drain().rev() {
+    for i in vec.drain(..).rev() {
         vec2.push(i);
     }
     assert_eq!(vec, []);
@@ -482,13 +482,43 @@ fn test_drain_items_reverse() {
 fn test_drain_items_zero_sized() {
     let mut vec = vec![(), (), ()];
     let mut vec2 = vec![];
-    for i in vec.drain() {
+    for i in vec.drain(..) {
         vec2.push(i);
     }
     assert_eq!(vec, []);
     assert_eq!(vec2, [(), (), ()]);
 }
 
+#[test]
+#[should_panic]
+fn drain_out_of_bounds_1() {
+    let mut v = vec![1, 2, 3, 4, 5];
+    v.drain(5..6);
+}
+
+#[test]
+fn drain_range() {
+    let mut v = vec![1, 2, 3, 4, 5];
+    for _ in v.drain(4..) {
+    }
+    assert_eq!(v, &[1, 2, 3, 4]);
+
+    let mut v: Vec<_> = (1..6).map(|x| x.to_string()).collect();
+    for _ in v.drain(1..4) {
+    }
+    assert_eq!(v, &[1.to_string(), 5.to_string()]);
+
+    let mut v: Vec<_> = (1..6).map(|x| x.to_string()).collect();
+    for _ in v.drain(1..4).rev() {
+    }
+    assert_eq!(v, &[1.to_string(), 5.to_string()]);
+
+    let mut v: Vec<_> = vec![(); 5];
+    for _ in v.drain(1..4).rev() {
+    }
+    assert_eq!(v, &[(), ()]);
+}
+
 #[test]
 fn test_into_boxed_slice() {
     let xs = vec![1, 2, 3];