Made vec_deque::Drain, hash_map::Drain, and hash_set::Drain covariant

5 files changed, 23 insertions, 9 deletions

diff --git a/src/libcollections/vec_deque.rs b/src/libcollections/vec_deque.rs
index 9c3792afa2f..812a67a7e78 100644
--- a/src/libcollections/vec_deque.rs
+++ b/src/libcollections/vec_deque.rs
@@ -24,6 +24,7 @@ use core::iter::{repeat, FromIterator};
 use core::mem;
 use core::ops::{Index, IndexMut};
 use core::ptr;
+use core::ptr::Shared;
 use core::slice;
 
 use core::hash::{Hash, Hasher};
@@ -903,7 +904,7 @@ impl<T> VecDeque<T> {
         self.head = drain_tail;
 
         Drain {
-            deque: self as *mut _,
+            deque: unsafe { Shared::new(self as *mut _) },
             after_tail: drain_head,
             after_head: head,
             iter: Iter {
@@ -1985,7 +1986,7 @@ pub struct Drain<'a, T: 'a> {
     after_tail: usize,
     after_head: usize,
     iter: Iter<'a, T>,
-    deque: *mut VecDeque<T>,
+    deque: Shared<VecDeque<T>>,
 }
 
 #[stable(feature = "drain", since = "1.6.0")]
@@ -1998,7 +1999,7 @@ impl<'a, T: 'a> Drop for Drain<'a, T> {
     fn drop(&mut self) {
         for _ in self.by_ref() {}
 
-        let source_deque = unsafe { &mut *self.deque };
+        let source_deque = unsafe { &mut **self.deque };
 
         // T = source_deque_tail; H = source_deque_head; t = drain_tail; h = drain_head
         //
diff --git a/src/libcollectionstest/vec_deque.rs b/src/libcollectionstest/vec_deque.rs
index a02666a50c2..5e8633a9748 100644
--- a/src/libcollectionstest/vec_deque.rs
+++ b/src/libcollectionstest/vec_deque.rs
@@ -10,6 +10,7 @@
 
 use std::collections::VecDeque;
 use std::fmt::Debug;
+use std::collections::vec_deque::Drain;
 
 use test;
 
@@ -999,3 +1000,8 @@ fn test_contains() {
 
     assert!(!v.contains(&3));
 }
+
+#[allow(dead_code)]
+fn assert_covariance() {
+    fn drain<'new>(d: Drain<'static, &'static str>) -> Drain<'new, &'new str> { d }
+}
diff --git a/src/libstd/collections/hash/map.rs b/src/libstd/collections/hash/map.rs
index fd7b0a2e6bb..e30cf638c6d 100644
--- a/src/libstd/collections/hash/map.rs
+++ b/src/libstd/collections/hash/map.rs
@@ -2036,6 +2036,8 @@ fn assert_covariance() {
     fn keys_val<'a, 'new>(v: Keys<'a, u8, &'static str>) -> Keys<'a, u8, &'new str> { v }
     fn values_key<'a, 'new>(v: Values<'a, &'static str, u8>) -> Values<'a, &'new str, u8> { v }
     fn values_val<'a, 'new>(v: Values<'a, u8, &'static str>) -> Values<'a, u8, &'new str> { v }
+    fn drain<'new>(d: Drain<'static, &'static str, &'static str>)
+        -> Drain<'new, &'new str, &'new str> { d }
 }
 
 #[cfg(test)]
diff --git a/src/libstd/collections/hash/set.rs b/src/libstd/collections/hash/set.rs
index e4ef3fca55d..837c5d8b8e1 100644
--- a/src/libstd/collections/hash/set.rs
+++ b/src/libstd/collections/hash/set.rs
@@ -1036,6 +1036,7 @@ fn assert_covariance() {
         -> Intersection<'a, &'new str, RandomState> { v }
     fn union<'a, 'new>(v: Union<'a, &'static str, RandomState>)
         -> Union<'a, &'new str, RandomState> { v }
+    fn drain<'new>(d: Drain<'static, &'static str>) -> Drain<'new, &'new str> { d }
 }
 
 #[cfg(test)]
diff --git a/src/libstd/collections/hash/table.rs b/src/libstd/collections/hash/table.rs
index c4c4cb45313..02931e5e389 100644
--- a/src/libstd/collections/hash/table.rs
+++ b/src/libstd/collections/hash/table.rs
@@ -17,7 +17,7 @@ use marker;
 use mem::{align_of, size_of};
 use mem;
 use ops::{Deref, DerefMut};
-use ptr::{self, Unique};
+use ptr::{self, Unique, Shared};
 
 use self::BucketState::*;
 
@@ -754,7 +754,8 @@ impl<K, V> RawTable<K, V> {
                 hashes_end: hashes_end,
                 marker: marker::PhantomData,
             },
-            table: self,
+            table: unsafe { Shared::new(self) },
+            marker: marker::PhantomData,
         }
     }
 
@@ -897,8 +898,9 @@ unsafe impl<K: Send, V: Send> Send for IntoIter<K, V> {}
 
 /// Iterator over the entries in a table, clearing the table.
 pub struct Drain<'a, K: 'a, V: 'a> {
-    table: &'a mut RawTable<K, V>,
+    table: Shared<RawTable<K, V>>,
     iter: RawBuckets<'static, K, V>,
+    marker: marker::PhantomData<&'a RawTable<K, V>>,
 }
 
 unsafe impl<'a, K: Sync, V: Sync> Sync for Drain<'a, K, V> {}
@@ -973,8 +975,8 @@ impl<'a, K, V> Iterator for Drain<'a, K, V> {
     #[inline]
     fn next(&mut self) -> Option<(SafeHash, K, V)> {
         self.iter.next().map(|bucket| {
-            self.table.size -= 1;
             unsafe {
+                (**self.table).size -= 1;
                 (SafeHash { hash: ptr::replace(bucket.hash, EMPTY_BUCKET) },
                  ptr::read(bucket.key),
                  ptr::read(bucket.val))
@@ -983,13 +985,15 @@ impl<'a, K, V> Iterator for Drain<'a, K, V> {
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let size = self.table.size();
+        let size = unsafe { (**self.table).size() };
         (size, Some(size))
     }
 }
 impl<'a, K, V> ExactSizeIterator for Drain<'a, K, V> {
     fn len(&self) -> usize {
-        self.table.size()
+        unsafe {
+            (**self.table).size()
+        }
     }
 }