From 2084f2ed77e28d78c8dfbd378251b29068a8c579 Mon Sep 17 00:00:00 2001
From: Corey Farwell <coreyf@rwell.org>
Date: Wed, 30 Mar 2016 19:12:29 -0400
Subject: [PATCH] Implement `values_mut` on `BTreeMap`.

https://github.com/rust-lang/rust/issues/32551
---
 src/libcollections/btree/map.rs     | 60 +++++++++++++++++++++++++++++
 src/libcollectionstest/btree/map.rs | 15 ++++++++
 src/libcollectionstest/lib.rs       |  1 +
 3 files changed, 76 insertions(+)

diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs
index 7819c7ef7969b..de40568fd6704 100644
--- a/src/libcollections/btree/map.rs
+++ b/src/libcollections/btree/map.rs
@@ -285,6 +285,12 @@ pub struct Values<'a, K: 'a, V: 'a> {
     inner: Iter<'a, K, V>,
 }
 
+/// A mutable iterator over a BTreeMap's values.
+#[unstable(feature = "map_values_mut", reason = "recently added", issue = "32551")]
+pub struct ValuesMut<'a, K: 'a, V: 'a> {
+    inner: IterMut<'a, K, V>,
+}
+
 /// An iterator over a sub-range of BTreeMap's entries.
 pub struct Range<'a, K: 'a, V: 'a> {
     front: Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
@@ -1006,6 +1012,33 @@ impl<'a, K, V> Iterator for Range<'a, K, V> {
     }
 }
 
+#[unstable(feature = "map_values_mut", reason = "recently added", issue = "32551")]
+impl<'a, K, V> Iterator for ValuesMut<'a, K, V> {
+    type Item = &'a mut V;
+
+    fn next(&mut self) -> Option<&'a mut V> {
+        self.inner.next().map(|(_, v)| v)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.inner.size_hint()
+    }
+}
+
+#[unstable(feature = "map_values_mut", reason = "recently added", issue = "32551")]
+impl<'a, K, V> DoubleEndedIterator for ValuesMut<'a, K, V> {
+    fn next_back(&mut self) -> Option<&'a mut V> {
+        self.inner.next_back().map(|(_, v)| v)
+    }
+}
+
+#[unstable(feature = "map_values_mut", reason = "recently added", issue = "32551")]
+impl<'a, K, V> ExactSizeIterator for ValuesMut<'a, K, V> {
+    fn len(&self) -> usize {
+        self.inner.len()
+    }
+}
+
 impl<'a, K, V> Range<'a, K, V> {
     unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
         let handle = self.front;
@@ -1403,6 +1436,33 @@ impl<K, V> BTreeMap<K, V> {
         Values { inner: self.iter() }
     }
 
+    /// Gets a mutable iterator over the values of the map, in order by key.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// # #![feature(map_values_mut)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut a = BTreeMap::new();
+    /// a.insert(1, String::from("hello"));
+    /// a.insert(2, String::from("goodbye"));
+    ///
+    /// for value in a.values_mut() {
+    ///     value.push_str("!");
+    /// }
+    ///
+    /// let values: Vec<String> = a.values().cloned().collect();
+    /// assert_eq!(values, [String::from("hello!"),
+    ///                     String::from("goodbye!")]);
+    /// ```
+    #[unstable(feature = "map_values_mut", reason = "recently added", issue = "32551")]
+    pub fn values_mut<'a>(&'a mut self) -> ValuesMut<'a, K, V> {
+        ValuesMut { inner: self.iter_mut() }
+    }
+
     /// Returns the number of elements in the map.
     ///
     /// # Examples
diff --git a/src/libcollectionstest/btree/map.rs b/src/libcollectionstest/btree/map.rs
index a27c2847a98bc..619bc189e6c09 100644
--- a/src/libcollectionstest/btree/map.rs
+++ b/src/libcollectionstest/btree/map.rs
@@ -114,6 +114,21 @@ fn test_iter_rev() {
     test(size, map.into_iter().rev());
 }
 
+#[test]
+fn test_values_mut() {
+    let mut a = BTreeMap::new();
+    a.insert(1, String::from("hello"));
+    a.insert(2, String::from("goodbye"));
+
+    for value in a.values_mut() {
+        value.push_str("!");
+    }
+
+    let values: Vec<String> = a.values().cloned().collect();
+    assert_eq!(values, [String::from("hello!"),
+                        String::from("goodbye!")]);
+}
+
 #[test]
 fn test_iter_mixed() {
     let size = 10000;
diff --git a/src/libcollectionstest/lib.rs b/src/libcollectionstest/lib.rs
index 62fefaa10f677..045ac1ce977fc 100644
--- a/src/libcollectionstest/lib.rs
+++ b/src/libcollectionstest/lib.rs
@@ -22,6 +22,7 @@
 #![feature(enumset)]
 #![feature(iter_arith)]
 #![feature(map_entry_keys)]
+#![feature(map_values_mut)]
 #![feature(pattern)]
 #![feature(rand)]
 #![feature(set_recovery)]