[collections] Adds `drain`: a way to sneak out the elements while clearing.

It is useful to move all the elements out of some collections without
deallocating the underlying buffer. It came up in IRC, and this patch
implements it as `drain`. This has been discussed as part of RFC 509.

r? @Gankro
cc: @frankmcsherry

3 files changed, 157 insertions, 10 deletions

diff --git a/src/libstd/collections/hash/map.rs b/src/libstd/collections/hash/map.rs
index 04dd5afdfa2..2020897fb3f 100644
--- a/src/libstd/collections/hash/map.rs
+++ b/src/libstd/collections/hash/map.rs
@@ -982,6 +982,35 @@ impl<K: Eq + Hash<S>, V, S, H: Hasher<S>> HashMap<K, V, H> {
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
     pub fn is_empty(&self) -> bool { self.len() == 0 }
 
+    /// Clears the map, returning all key-value pairs as an iterator. Keeps the
+    /// allocated memory for reuse.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use std::collections::HashMap;
+    ///
+    /// let mut a = HashMap::new();
+    /// a.insert(1u, "a");
+    /// a.insert(2u, "b");
+    ///
+    /// for (k, v) in a.drain().take(1) {
+    ///     assert!(k == 1 || k == 2);
+    ///     assert!(v == "a" || v == "b");
+    /// }
+    ///
+    /// assert!(a.is_empty());
+    /// ```
+    #[inline]
+    #[unstable = "matches collection reform specification, waiting for dust to settle"]
+    pub fn drain(&mut self) -> Drain<K, V> {
+        fn last_two<A, B, C>((_, b, c): (A, B, C)) -> (B, C) { (b, c) }
+
+        Drain {
+            inner: self.table.drain().map(last_two),
+        }
+    }
+
     /// Clears the map, removing all key-value pairs. Keeps the allocated memory
     /// for reuse.
     ///
@@ -996,16 +1025,9 @@ impl<K: Eq + Hash<S>, V, S, H: Hasher<S>> HashMap<K, V, H> {
     /// assert!(a.is_empty());
     /// ```
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
+    #[inline]
     pub fn clear(&mut self) {
-        let cap = self.table.capacity();
-        let mut buckets = Bucket::first(&mut self.table);
-
-        while buckets.index() != cap {
-            buckets = match buckets.peek() {
-                Empty(b)   => b.next(),
-                Full(full) => full.take().0.next(),
-            };
-        }
+        self.drain();
     }
 
     /// Deprecated: Renamed to `get`.
@@ -1306,6 +1328,16 @@ pub struct Values<'a, K: 'a, V: 'a> {
     inner: Map<(&'a K, &'a V), &'a V, Entries<'a, K, V>, fn((&'a K, &'a V)) -> &'a V>
 }
 
+/// HashMap drain iterator
+pub struct Drain<'a, K: 'a, V: 'a> {
+    inner: iter::Map<
+        (SafeHash, K, V),
+        (K, V),
+        table::Drain<'a, K, V>,
+        fn((SafeHash, K, V)) -> (K, V),
+    >
+}
+
 /// A view into a single occupied location in a HashMap
 pub struct OccupiedEntry<'a, K:'a, V:'a> {
     elem: FullBucket<K, V, &'a mut RawTable<K, V>>,
@@ -1360,6 +1392,17 @@ impl<'a, K, V> Iterator<&'a V> for Values<'a, K, V> {
     #[inline] fn size_hint(&self) -> (uint, Option<uint>) { self.inner.size_hint() }
 }
 
+impl<'a, K: 'a, V: 'a> Iterator<(K, V)> for Drain<'a, K, V> {
+    #[inline]
+    fn next(&mut self) -> Option<(K, V)> {
+        self.inner.next()
+    }
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        self.inner.size_hint()
+    }
+}
+
 impl<'a, K, V> OccupiedEntry<'a, K, V> {
     /// Gets a reference to the value in the entry
     pub fn get(&self) -> &V {
diff --git a/src/libstd/collections/hash/set.rs b/src/libstd/collections/hash/set.rs
index 67c0f887832..d75602a2716 100644
--- a/src/libstd/collections/hash/set.rs
+++ b/src/libstd/collections/hash/set.rs
@@ -21,7 +21,7 @@ use iter::{Iterator, IteratorExt, FromIterator, Map, FilterMap, Chain, Repeat, Z
 use option::Option::{Some, None, mod};
 use result::Result::{Ok, Err};
 
-use super::map::{HashMap, MoveEntries, Keys, INITIAL_CAPACITY};
+use super::map::{mod, HashMap, MoveEntries, Keys, INITIAL_CAPACITY};
 
 // FIXME(conventions): implement BitOr, BitAnd, BitXor, and Sub
 
@@ -420,6 +420,14 @@ impl<T: Eq + Hash<S>, S, H: Hasher<S>> HashSet<T, H> {
     #[unstable = "matches collection reform specification, waiting for dust to settle"]
     pub fn is_empty(&self) -> bool { self.map.len() == 0 }
 
+    /// Clears the set, returning all elements in an iterator.
+    #[inline]
+    #[unstable = "matches collection reform specification, waiting for dust to settle"]
+    pub fn drain(&mut self) -> Drain<T> {
+        fn first<A, B>((a, _): (A, B)) -> A { a }
+        Drain { iter: self.map.drain().map(first) }
+    }
+
     /// Clears the set, removing all values.
     ///
     /// # Example
@@ -626,6 +634,11 @@ pub struct SetMoveItems<K> {
     iter: Map<(K, ()), K, MoveEntries<K, ()>, fn((K, ())) -> K>
 }
 
+/// HashSet drain iterator
+pub struct Drain<'a, K: 'a> {
+    iter: Map<(K, ()), K, map::Drain<'a, K, ()>, fn((K, ())) -> K>,
+}
+
 // `Repeat` is used to feed the filter closure an explicit capture
 // of a reference to the other set
 /// Set operations iterator, used directly for intersection and difference
@@ -658,6 +671,11 @@ impl<K> Iterator<K> for SetMoveItems<K> {
     fn size_hint(&self) -> (uint, Option<uint>) { self.iter.size_hint() }
 }
 
+impl<'a, K: 'a> Iterator<K> for Drain<'a, K> {
+    fn next(&mut self) -> Option<K> { self.iter.next() }
+    fn size_hint(&self) -> (uint, Option<uint>) { self.iter.size_hint() }
+}
+
 impl<'a, T, H> Iterator<&'a T> for SetAlgebraItems<'a, T, H> {
     fn next(&mut self) -> Option<&'a T> { self.iter.next() }
     fn size_hint(&self) -> (uint, Option<uint>) { self.iter.size_hint() }
@@ -914,4 +932,41 @@ mod test_set {
         assert!(set_str == "{1, 2}" || set_str == "{2, 1}");
         assert_eq!(format!("{}", empty), "{}");
     }
+
+    #[test]
+    fn test_trivial_drain() {
+        let mut s = HashSet::<int>::new();
+        for _ in s.drain() {}
+        assert!(s.is_empty());
+        drop(s);
+
+        let mut s = HashSet::<int>::new();
+        drop(s.drain());
+        assert!(s.is_empty());
+    }
+
+    #[test]
+    fn test_drain() {
+        let mut s: HashSet<int> = range(1, 100).collect();
+
+        // try this a bunch of times to make sure we don't screw up internal state.
+        for _ in range(0i, 20) {
+            assert_eq!(s.len(), 99);
+
+            {
+                let mut last_i = 0;
+                let mut d = s.drain();
+                for (i, x) in d.by_ref().take(50).enumerate() {
+                    last_i = i;
+                    assert!(x != 0);
+                }
+                assert_eq!(last_i, 49);
+            }
+
+            for _ in s.iter() { panic!("s should be empty!"); }
+
+            // reset to try again.
+            s.extend(range(1, 100));
+        }
+    }
 }
diff --git a/src/libstd/collections/hash/table.rs b/src/libstd/collections/hash/table.rs
index da06387e9a5..115edcabca1 100644
--- a/src/libstd/collections/hash/table.rs
+++ b/src/libstd/collections/hash/table.rs
@@ -684,6 +684,19 @@ impl<K, V> RawTable<K, V> {
         }
     }
 
+    pub fn drain(&mut self) -> Drain<K, V> {
+        let RawBuckets { raw, hashes_end, .. } = self.raw_buckets();
+        // Replace the marker regardless of lifetime bounds on parameters.
+        Drain {
+            iter: RawBuckets {
+                raw: raw,
+                hashes_end: hashes_end,
+                marker: marker::ContravariantLifetime::<'static>,
+            },
+            table: self,
+        }
+    }
+
     /// Returns an iterator that copies out each entry. Used while the table
     /// is being dropped.
     unsafe fn rev_move_buckets(&mut self) -> RevMoveBuckets<K, V> {
@@ -774,6 +787,12 @@ pub struct MoveEntries<K, V> {
     iter: RawBuckets<'static, 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>,
+    iter: RawBuckets<'static, K, V>,
+}
+
 impl<'a, K, V> Iterator<(&'a K, &'a V)> for Entries<'a, K, V> {
     fn next(&mut self) -> Option<(&'a K, &'a V)> {
         self.iter.next().map(|bucket| {
@@ -828,6 +847,36 @@ impl<K, V> Iterator<(SafeHash, K, V)> for MoveEntries<K, V> {
     }
 }
 
+impl<'a, K: 'a, V: 'a> Iterator<(SafeHash, K, V)> for Drain<'a, K, V> {
+    #[inline]
+    fn next(&mut self) -> Option<(SafeHash, K, V)> {
+        self.iter.next().map(|bucket| {
+            self.table.size -= 1;
+            unsafe {
+                (
+                    SafeHash {
+                        hash: ptr::replace(bucket.hash, EMPTY_BUCKET),
+                    },
+                    ptr::read(bucket.key as *const K),
+                    ptr::read(bucket.val as *const V)
+                )
+            }
+        })
+    }
+
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let size = self.table.size();
+        (size, Some(size))
+    }
+}
+
+#[unsafe_destructor]
+impl<'a, K: 'a, V: 'a> Drop for Drain<'a, K, V> {
+    fn drop(&mut self) {
+        for _ in *self {}
+    }
+}
+
 impl<K: Clone, V: Clone> Clone for RawTable<K, V> {
     fn clone(&self) -> RawTable<K, V> {
         unsafe {