Make vec::Drain and binary_heap::Drain covariant

Author: tomgarcia

src/libcollections/binary_heap.rs

@@ -1016,12 +1016,12 @@ impl<T> ExactSizeIterator for IntoIter<T> {}
 
 /// An iterator that drains a `BinaryHeap`.
 #[stable(feature = "drain", since = "1.6.0")]
-pub struct Drain<'a, T: 'a> {
-    iter: vec::Drain<'a, T>,
+pub struct Drain<T> {
+    iter: vec::Drain<T>,
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T: 'a> Iterator for Drain<'a, T> {
+impl<T> Iterator for Drain<T> {
     type Item = T;
 
     #[inline]
@@ -1036,15 +1036,15 @@ impl<'a, T: 'a> Iterator for Drain<'a, T> {
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T: 'a> DoubleEndedIterator for Drain<'a, T> {
+impl<T> DoubleEndedIterator for Drain<T> {
     #[inline]
     fn next_back(&mut self) -> Option<T> {
         self.iter.next_back()
     }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T: 'a> ExactSizeIterator for Drain<'a, T> {}
+impl<T> ExactSizeIterator for Drain<T> {}
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T: Ord> From<Vec<T>> for BinaryHeap<T> {

src/libcollections/vec.rs

@@ -74,6 +74,7 @@ use core::ops::{Index, IndexMut};
 use core::ops;
 use core::ptr;
 use core::slice;
+use super::vec_deque::VecDeque;
 
 use super::SpecExtend;
 use super::range::RangeArgument;
@@ -843,20 +844,20 @@ impl<T> Vec<T> {
         let end = *range.end().unwrap_or(&len);
         assert!(start <= end);
         assert!(end <= len);
+        let mut drain_vec = VecDeque::new();
 
         unsafe {
-            // 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 {
-                tail_start: end,
-                tail_len: len - end,
-                iter: range_slice.iter_mut(),
-                vec: self as *mut _,
+            for i in start..end {
+                    let p = self.as_ptr().offset(i as isize);
+                    drain_vec.push_back(ptr::read(p));
             }
+            let src = self.as_ptr().offset(end as isize);
+            let dst = self.as_mut_ptr().offset(start as isize);
+            ptr::copy(src, dst, len - end);
+            self.set_len(len - (end - start));
+        }
+        Drain {
+            deque: drain_vec
         }
     }
 
@@ -1755,64 +1756,43 @@ impl<T> Drop for IntoIter<T> {
 /// [`drain`]: struct.Vec.html#method.drain
 /// [`Vec`]: struct.Vec.html
 #[stable(feature = "drain", since = "1.6.0")]
-pub struct Drain<'a, T: 'a> {
-    /// Index of tail to preserve
-    tail_start: usize,
-    /// Length of tail
-    tail_len: usize,
+pub struct Drain<T> {
     /// Current remaining range to remove
-    iter: slice::IterMut<'a, T>,
-    vec: *mut Vec<T>,
+    deque: VecDeque<T>
 }
 
 #[stable(feature = "drain", since = "1.6.0")]
-unsafe impl<'a, T: Sync> Sync for Drain<'a, T> {}
+unsafe impl<T: Sync> Sync for Drain<T> {}
 #[stable(feature = "drain", since = "1.6.0")]
-unsafe impl<'a, T: Send> Send for Drain<'a, T> {}
+unsafe impl<T: Send> Send for Drain<T> {}
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> Iterator for Drain<'a, T> {
+impl<T> Iterator for Drain<T> {
     type Item = T;
 
     #[inline]
     fn next(&mut self) -> Option<T> {
-        self.iter.next().map(|elt| unsafe { ptr::read(elt as *const _) })
+        self.deque.pop_front()
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        self.iter.size_hint()
+        (self.deque.len(), Some(self.deque.len()))
     }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> DoubleEndedIterator for Drain<'a, T> {
+impl<T> DoubleEndedIterator for Drain<T> {
     #[inline]
     fn next_back(&mut self) -> Option<T> {
-        self.iter.next_back().map(|elt| unsafe { ptr::read(elt as *const _) })
+        self.deque.pop_back()
     }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<'a, T> Drop for Drain<'a, T> {
+impl<T> Drop for Drain<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> {}
+impl<T> ExactSizeIterator for Drain<T> {}

src/libcollectionstest/binary_heap.rs

@@ -9,6 +9,7 @@
 // except according to those terms.
 
 use std::collections::BinaryHeap;
+use std::collections::binary_heap::Drain;
 
 #[test]
 fn test_iterator() {
@@ -292,3 +293,8 @@ fn test_extend_specialization() {
 
     assert_eq!(a.into_sorted_vec(), [-20, -10, 1, 2, 3, 3, 5, 43]);
 }
+
+#[allow(dead_code)]
+fn assert_covariance() {
+    fn drain<'new>(d: Drain<&'static str>) -> Drain<&'new str> { d }
+}

src/libcollectionstest/vec.rs

@@ -11,6 +11,7 @@
 use std::borrow::Cow;
 use std::iter::{FromIterator, repeat};
 use std::mem::size_of;
+use std::vec::Drain;
 
 use test::Bencher;
 
@@ -510,6 +511,11 @@ fn test_cow_from() {
     }
 }
 
+#[allow(dead_code)]
+fn assert_covariance() {
+    fn drain<'new>(d: Drain<&'static str>) -> Drain<&'new str> { d }
+}
+
 #[bench]
 fn bench_new(b: &mut Bencher) {
     b.iter(|| {