std: add consuming iterators for `HashMap` and `HashSet`

diff --git a/src/libstd/hashmap.rs b/src/libstd/hashmap.rs
index 0e31e47d56b2742548356b3759dd7b708180a216..3b3b71d729735281acfb91ddf9479a4ba93e7538 100644 (file)
--- a/src/libstd/hashmap.rs
+++ b/src/libstd/hashmap.rs
@@ -455,6 +455,14 @@ pub fn consume(&mut self, f: &fn(K, V)) {
         }
     }
 
+    /// Creates a consuming iterator, that is, one that moves each key-value
+    /// pair out of the map in arbitrary order. The map cannot be used after
+    /// calling this.
+    pub fn consume_iter(self) -> HashMapConsumeIterator<K, V> {
+        // `consume_rev_iter` is more efficient than `consume_iter` for vectors
+        HashMapConsumeIterator {iter: self.buckets.consume_rev_iter()}
+    }
+
     /// Retrieves a value for the given key, failing if the key is not
     /// present.
     pub fn get<'a>(&'a self, k: &K) -> &'a V {
@@ -568,11 +576,21 @@ pub struct HashMapMutIterator<'self, K, V> {
     priv iter: vec::VecMutIterator<'self, Option<Bucket<K, V>>>,
 }
 
+/// HashMap consume iterator
+pub struct HashMapConsumeIterator<K, V> {
+    priv iter: vec::VecConsumeRevIterator<Option<Bucket<K, V>>>,
+}
+
 /// HashSet iterator
 pub struct HashSetIterator<'self, K> {
     priv iter: vec::VecIterator<'self, Option<Bucket<K, ()>>>,
 }
 
+/// HashSet consume iterator
+pub struct HashSetConsumeIterator<K> {
+    priv iter: vec::VecConsumeRevIterator<Option<Bucket<K, ()>>>,
+}
+
 impl<'self, K, V> Iterator<(&'self K, &'self V)> for HashMapIterator<'self, K, V> {
     #[inline]
     fn next(&mut self) -> Option<(&'self K, &'self V)> {
@@ -599,6 +617,19 @@ fn next(&mut self) -> Option<(&'self K, &'self mut V)> {
     }
 }
 
+impl<K, V> Iterator<(K, V)> for HashMapConsumeIterator<K, V> {
+    #[inline]
+    fn next(&mut self) -> Option<(K, V)> {
+        for self.iter.advance |elt| {
+            match elt {
+                Some(Bucket {key, value, _}) => return Some((key, value)),
+                None => {},
+            }
+        }
+        None
+    }
+}
+
 impl<'self, K> Iterator<&'self K> for HashSetIterator<'self, K> {
     #[inline]
     fn next(&mut self) -> Option<&'self K> {
@@ -612,6 +643,19 @@ fn next(&mut self) -> Option<&'self K> {
     }
 }
 
+impl<K> Iterator<K> for HashSetConsumeIterator<K> {
+    #[inline]
+    fn next(&mut self) -> Option<K> {
+        for self.iter.advance |elt| {
+            match elt {
+                Some(bucket) => return Some(bucket.key),
+                None => {},
+            }
+        }
+        None
+    }
+}
+
 impl<K: Eq + Hash, V, T: Iterator<(K, V)>> FromIterator<(K, V), T> for HashMap<K, V> {
     pub fn from_iterator(iter: &mut T) -> HashMap<K, V> {
         let (lower, _) = iter.size_hint();
@@ -726,6 +770,14 @@ pub fn consume(&mut self, f: &fn(T)) {
         self.map.consume(|k, _| f(k))
     }
 
+    /// Creates a consuming iterator, that is, one that moves each value out
+    /// of the set in arbitrary order. The set cannot be used after calling
+    /// this.
+    pub fn consume_iter(self) -> HashSetConsumeIterator<T> {
+        // `consume_rev_iter` is more efficient than `consume_iter` for vectors
+        HashSetConsumeIterator {iter: self.map.buckets.consume_rev_iter()}
+    }
+
     /// Returns true if the hash set contains a value equivalent to the
     /// given query value.
     pub fn contains_equiv<Q:Hash + Equiv<T>>(&self, value: &Q) -> bool {
@@ -888,6 +940,21 @@ fn test_consume_still_usable() {
         assert!(m.insert(1, 2));
     }
 
+    #[test]
+    fn test_consume_iter() {
+        let hm = {
+            let mut hm = HashMap::new();
+
+            hm.insert('a', 1);
+            hm.insert('b', 2);
+
+            hm
+        };
+
+        let v = hm.consume_iter().collect::<~[(char, int)]>();
+        assert!([('a', 1), ('b', 2)] == v || [('b', 2), ('a', 1)] == v);
+    }
+
     #[test]
     fn test_iterate() {
         let mut m = linear_map_with_capacity(4);
@@ -1168,4 +1235,19 @@ fn test_from_iter() {
             assert!(set.contains(x));
         }
     }
+
+    #[test]
+    fn test_consume_iter() {
+        let hs = {
+            let mut hs = HashSet::new();
+
+            hs.insert('a');
+            hs.insert('b');
+
+            hs
+        };
+
+        let v = hs.consume_iter().collect::<~[char]>();
+        assert!(['a', 'b'] == v || ['b', 'a'] == v);
+    }
 }