Auto merge of #24781 - bluss:vec-drain-range, r=alexcrichton

Implement Vec::drain(\<range type\>) from rust-lang/rfcs#574, tracking issue #23055.

This is a big step forward for vector usability. This is an introduction of an API for removing a range of *m* consecutive elements from a vector, as efficently as possible.

New features:

- Introduce trait `std::collections::range::RangeArgument` implemented by all four built-in range types.
- Change `Vec::drain()` to use `Vec::drain<R: RangeArgument>(R)`

Implementation notes:

- Use @gankro's idea for memory safety: Use `set_len` on the source vector when creating the iterator, to make sure that the part of the vector that will be modified is unreachable. Fix up things in Drain's destructor — but even if it doesn't run, we don't expose any moved-out-from slots of the vector.
- This `.drain<R>(R)` very close to how it is specified in the RFC.
- Introduced as unstable
- Drain reuses the slice iterator — copying and pasting the same iterator pointer arithmetic again felt very bad
- The `usize` index as a range argument in the RFC is not included. The ranges trait would have to change to accomodate it.

Please help me with:

- Name and location of the new ranges trait.
- Design of the ranges trait
- Understanding Niko's comments about variance (Note: for a long time I was using a straight up &mut Vec in the iterator, but I changed this to permit reusing the slice iterator).

Previous PR and discussion: #23071

Author: bors

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/lib.rs

@@ -42,7 +42,8 @@
 #![feature(slice_patterns)]
 #![feature(debug_builders)]
 #![feature(utf8_error)]
-#![cfg_attr(test, feature(rand, rustc_private, test, hash, collections))]
+#![cfg_attr(test, feature(rand, rustc_private, test, hash, collections,
+                          collections_drain, collections_range))]
 #![cfg_attr(test, allow(deprecated))] // rand
 
 #![feature(no_std)]
@@ -82,6 +83,7 @@ pub mod borrow;
 pub mod enum_set;
 pub mod fmt;
 pub mod linked_list;
+pub mod range;
 pub mod slice;
 pub mod str;
 pub mod string;

src/libcollections/range.rs

@@ -0,0 +1,45 @@
+// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// 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.
+#![unstable(feature = "collections_range", reason = "was just added")]
+
+//! Range syntax.
+
+use core::option::Option::{self, None, Some};
+use core::ops::{RangeFull, Range, RangeTo, RangeFrom};
+
+/// **RangeArgument** is implemented by Rust's built-in range types, produced
+/// by range syntax like `..`, `a..`, `..b` or `c..d`.
+pub trait RangeArgument<T> {
+    /// Start index (inclusive)
+    ///
+    /// Return start value if present, else `None`.
+    fn start(&self) -> Option<&T> { None }
+
+    /// End index (exclusive)
+    ///
+    /// Return end value if present, else `None`.
+    fn end(&self) -> Option<&T> { None }
+}
+
+
+impl<T> RangeArgument<T> for RangeFull {}
+
+impl<T> RangeArgument<T> for RangeFrom<T> {
+    fn start(&self) -> Option<&T> { Some(&self.start) }
+}
+
+impl<T> RangeArgument<T> for RangeTo<T> {
+    fn end(&self) -> Option<&T> { Some(&self.end) }
+}
+
+impl<T> RangeArgument<T> for Range<T> {
+    fn start(&self) -> Option<&T> { Some(&self.start) }
+    fn end(&self) -> Option<&T> { Some(&self.end) }
+}

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;
 
@@ -714,36 +716,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());
+    /// # #![feature(collections_drain, collections_range)]
+    ///
+    /// // 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> {
+    #[unstable(feature = "collections_drain",
+               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);
+            // Use the borrow in the IterMut to indicate borrowing behavior of the
+            // whole Drain iterator (like &mut T).
+            let range_slice = slice::from_raw_parts_mut(
+                                        self.as_mut_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_mut(),
+                vec: self as *mut _,
             }
         }
     }
@@ -1795,14 +1822,16 @@ 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_drain", 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::IterMut<'a, T>,
+    vec: *mut Vec<T>,
 }
 
 unsafe impl<'a, T: Sync> Sync for Drain<'a, T> {}
@@ -1814,76 +1843,56 @@ impl<'a, T> Iterator for Drain<'a, 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))
-                }
+        self.iter.next().map(|elt|
+            unsafe {
+                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))
+        self.iter.size_hint()
     }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 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))
-                }
+        self.iter.next_back().map(|elt|
+            unsafe {
+                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].
+        // exhaust self first
+        while let Some(_) = self.next() { }
 
-        // destroy the remaining elements
-        for _x in self.by_ref() {}
+        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/lib.rs

@@ -10,6 +10,7 @@
 
 #![feature(box_syntax)]
 #![feature(collections)]
+#![feature(collections_drain)]
 #![feature(core)]
 #![feature(hash)]
 #![feature(rand)]

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 test_drain_out_of_bounds() {
+    let mut v = vec![1, 2, 3, 4, 5];
+    v.drain(5..6);
+}
+
+#[test]
+fn test_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];

src/test/run-pass/sync-send-iterators-in-libcollections.rs

@@ -12,6 +12,7 @@
 
 #![allow(unused_mut)]
 #![feature(collections)]
+#![feature(collections_drain)]
 
 extern crate collections;
 
@@ -96,5 +97,6 @@ fn main() {
 
     all_sync_send!(VecMap::<usize>::new(), iter, iter_mut, drain, into_iter, keys, values);
 
-    all_sync_send!(Vec::<usize>::new(), into_iter, drain);
+    all_sync_send!(Vec::<usize>::new(), into_iter);
+    is_sync_send!(Vec::<usize>::new(), drain(..));
 }