rollup merge of #19993: bluss/setalgebraitems

[rust.git] / src / libstd / collections / hash / set.rs

diff --git a/src/libstd/collections/hash/set.rs b/src/libstd/collections/hash/set.rs
index 99fe6696ec78d92157169245f10142608c1fa763..f587669d3dacd5e967aa9d9e38e97a4b1ff4afa9 100644 (file)
--- a/src/libstd/collections/hash/set.rs
+++ b/src/libstd/collections/hash/set.rs
@@ -21,7 +21,7 @@
 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
 
@@ -414,6 +414,14 @@ pub fn len(&self) -> uint { self.map.len() }
     #[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
@@ -620,6 +628,11 @@ pub struct IntoIter<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>,
+}
+
 /// Intersection iterator
 pub struct Intersection<'a, T: 'a, H: 'a> {
     // iterator of the first set
@@ -656,6 +669,11 @@ fn next(&mut self) -> Option<K> { self.iter.next() }
     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, S, H> Iterator<&'a T> for Intersection<'a, T, H>
     where T: Eq + Hash<S>, H: Hasher<S>
 {
@@ -715,7 +733,6 @@ mod test_set {
     use prelude::*;
 
     use super::HashSet;
-    use slice::PartialEqSliceExt;
 
     #[test]
     fn test_disjoint() {
@@ -951,4 +968,41 @@ fn test_show() {
         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));
+        }
+    }
 }