Merge pull request #7679 from alexcrichton/consume-smallintmap

Add a `consume` method to SmallIntMap

1 files changed, 28 insertions, 1 deletions

diff --git a/src/libextra/smallintmap.rs b/src/libextra/smallintmap.rs
index d952374ee5c..d79a8cc426a 100644
--- a/src/libextra/smallintmap.rs
+++ b/src/libextra/smallintmap.rs
@@ -18,10 +18,11 @@
 
 use std::cmp;
 use std::container::{Container, Mutable, Map, Set};
-use std::iterator::{Iterator,IteratorUtil,ZipIterator,Counter};
+use std::iterator::*;
 use std::uint;
 use std::util::replace;
 use std::vec::{VecIterator,VecMutIterator,VecRevIterator,VecMutRevIterator};
+use std::vec::VecConsumeIterator;
 
 #[allow(missing_doc)]
 pub struct SmallIntMap<T> {
@@ -204,6 +205,17 @@ impl<V> SmallIntMap<V> {
             iter: Counter::new(self.len() as int - 1, -1).zip(self.v.mut_rev_iter())
         }
     }
+
+    /// Empties the hash map, moving all values into the specified closure
+    pub fn consume(&mut self)
+        -> FilterMapIterator<(uint, Option<V>), (uint, V),
+                EnumerateIterator<Option<V>, VecConsumeIterator<Option<V>>>>
+    {
+        let values = replace(&mut self.v, ~[]);
+        values.consume_iter().enumerate().filter_map(|(i, v)| {
+            v.map_consume(|v| (i, v))
+        })
+    }
 }
 
 impl<V:Copy> SmallIntMap<V> {
@@ -625,6 +637,21 @@ mod tests {
 
         assert!(a.iter().all(|(_,v)| *v == 5 ));
     }
+
+    #[test]
+    fn test_consume() {
+        let mut m = SmallIntMap::new();
+        m.insert(1, ~2);
+        let mut called = false;
+        for m.consume().advance |(k, v)| {
+            assert!(!called);
+            called = true;
+            assert_eq!(k, 1);
+            assert_eq!(v, ~2);
+        }
+        assert!(called);
+        m.insert(2, ~1);
+    }
 }
 
 #[cfg(test)]