From c3ac929ba966666fbe9c72eb97e38d5e37ad11e2 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Adolfo=20Ochagav=C3=ADa?= <aochagavia92@gmail.com>
Date: Mon, 19 Jan 2015 11:06:15 +0100
Subject: [PATCH] Add a Drain iterator to VecMap

---
 src/libcollections/vec_map.rs | 61 +++++++++++++++++++++++++++++++----
 1 file changed, 55 insertions(+), 6 deletions(-)

diff --git a/src/libcollections/vec_map.rs b/src/libcollections/vec_map.rs
index 94ae1bece86..93f3e192d6d 100644
--- a/src/libcollections/vec_map.rs
+++ b/src/libcollections/vec_map.rs
@@ -186,7 +186,7 @@ pub fn reserve_len_exact(&mut self, len: uint) {
         }
     }
 
-    /// Returns an iterator visiting all keys in ascending order by the keys.
+    /// Returns an iterator visiting all keys in ascending order of the keys.
     /// The iterator's element type is `uint`.
     #[stable]
     pub fn keys<'r>(&'r self) -> Keys<'r, V> {
@@ -196,7 +196,7 @@ fn first<A, B>((a, _): (A, B)) -> A { a }
         Keys { iter: self.iter().map(first) }
     }
 
-    /// Returns an iterator visiting all values in ascending order by the keys.
+    /// Returns an iterator visiting all values in ascending order of the keys.
     /// The iterator's element type is `&'r V`.
     #[stable]
     pub fn values<'r>(&'r self) -> Values<'r, V> {
@@ -206,7 +206,7 @@ fn second<A, B>((_, b): (A, B)) -> B { b }
         Values { iter: self.iter().map(second) }
     }
 
-    /// Returns an iterator visiting all key-value pairs in ascending order by the keys.
+    /// Returns an iterator visiting all key-value pairs in ascending order of the keys.
     /// The iterator's element type is `(uint, &'r V)`.
     ///
     /// # Examples
@@ -233,7 +233,7 @@ pub fn iter<'r>(&'r self) -> Iter<'r, V> {
         }
     }
 
-    /// Returns an iterator visiting all key-value pairs in ascending order by the keys,
+    /// Returns an iterator visiting all key-value pairs in ascending order of the keys,
     /// with mutable references to the values.
     /// The iterator's element type is `(uint, &'r mut V)`.
     ///
@@ -264,7 +264,7 @@ pub fn iter_mut<'r>(&'r mut self) -> IterMut<'r, V> {
         }
     }
 
-    /// Returns an iterator visiting all key-value pairs in ascending order by
+    /// Returns an iterator visiting all key-value pairs in ascending order of
     /// the keys, consuming the original `VecMap`.
     /// The iterator's element type is `(uint, &'r V)`.
     ///
@@ -278,7 +278,6 @@ pub fn iter_mut<'r>(&'r mut self) -> IterMut<'r, V> {
     /// map.insert(3, "c");
     /// map.insert(2, "b");
     ///
-    /// // Not possible with .iter()
     /// let vec: Vec<(uint, &str)> = map.into_iter().collect();
     ///
     /// assert_eq!(vec, vec![(1, "a"), (2, "b"), (3, "c")]);
@@ -293,6 +292,34 @@ fn filter<A>((i, v): (uint, Option<A>)) -> Option<(uint, A)> {
         IntoIter { iter: self.v.into_iter().enumerate().filter_map(filter) }
     }
 
+    /// Returns an iterator visiting all key-value pairs in ascending order of
+    /// the keys, emptying (but not consuming) the original `VecMap`.
+    /// The iterator's element type is `(uint, &'r V)`. Keeps the allocated memory for reuse.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::collections::VecMap;
+    ///
+    /// let mut map = VecMap::new();
+    /// map.insert(1, "a");
+    /// map.insert(3, "c");
+    /// map.insert(2, "b");
+    ///
+    /// let vec: Vec<(uint, &str)> = map.drain().collect();
+    ///
+    /// assert_eq!(vec, vec![(1, "a"), (2, "b"), (3, "c")]);
+    /// ```
+    #[unstable = "matches collection reform specification, waiting for dust to settle"]
+    pub fn drain<'a>(&'a mut self) -> Drain<'a, V> {
+        fn filter<A>((i, v): (uint, Option<A>)) -> Option<(uint, A)> {
+            v.map(|v| (i, v))
+        }
+        let filter: fn((uint, Option<V>)) -> Option<(uint, V)> = filter; // coerce to fn ptr
+
+        Drain { iter: self.v.drain().enumerate().filter_map(filter) }
+    }
+
     /// Return the number of elements in the map.
     ///
     /// # Examples
@@ -672,6 +699,28 @@ pub struct IntoIter<V> {
     fn((uint, Option<V>)) -> Option<(uint, V)>>
 }
 
+#[unstable]
+pub struct Drain<'a, V> {
+    iter: FilterMap<
+    (uint, Option<V>),
+    (uint, V),
+    Enumerate<vec::Drain<'a, Option<V>>>,
+    fn((uint, Option<V>)) -> Option<(uint, V)>>
+}
+
+#[unstable]
+impl<'a, V> Iterator for Drain<'a, V> {
+    type Item = (uint, V);
+
+    fn next(&mut self) -> Option<(uint, V)> { self.iter.next() }
+    fn size_hint(&self) -> (uint, Option<uint>) { self.iter.size_hint() }
+}
+
+#[unstable]
+impl<'a, V> DoubleEndedIterator for Drain<'a, V> {
+    fn next_back(&mut self) -> Option<(uint, V)> { self.iter.next_back() }
+}
+
 #[stable]
 impl<'a, V> Iterator for Keys<'a, V> {
     type Item = uint;
-- 
2.44.0