Auto merge of #32633 - frewsxcv:map-values-mut, r=alexcrichton

Implement `values_mut` on `{BTree, Hash}Map`

https://github.com/rust-lang/rust/issues/32551

5 files changed, 138 insertions, 0 deletions

diff --git a/src/libcollections/btree/map.rs b/src/libcollections/btree/map.rs
index 7819c7ef796..de40568fd67 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 a27c2847a98..619bc189e6c 100644
--- a/src/libcollectionstest/btree/map.rs
+++ b/src/libcollectionstest/btree/map.rs
@@ -115,6 +115,21 @@ fn test_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 62fefaa10f6..045ac1ce977 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)]
diff --git a/src/libstd/collections/hash/map.rs b/src/libstd/collections/hash/map.rs
index 80b5448800e..4c29023660a 100644
--- a/src/libstd/collections/hash/map.rs
+++ b/src/libstd/collections/hash/map.rs
@@ -861,6 +861,34 @@ impl<K, V, S> HashMap<K, V, S>
         Values { inner: self.iter() }
     }
 
+    /// An iterator visiting all values mutably in arbitrary order.
+    /// Iterator element type is `&'a mut V`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #![feature(map_values_mut)]
+    /// use std::collections::HashMap;
+    ///
+    /// let mut map = HashMap::new();
+    ///
+    /// map.insert("a", 1);
+    /// map.insert("b", 2);
+    /// map.insert("c", 3);
+    ///
+    /// for val in map.values_mut() {
+    ///     *val = *val + 10;
+    /// }
+    ///
+    /// for val in map.values() {
+    ///     print!("{}", val);
+    /// }
+    /// ```
+    #[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() }
+    }
+
     /// An iterator visiting all key-value pairs in arbitrary order.
     /// Iterator element type is `(&'a K, &'a V)`.
     ///
@@ -1262,6 +1290,12 @@ pub struct Drain<'a, K: 'a, V: 'a> {
     inner: table::Drain<'a, K, V>
 }
 
+/// Mutable HashMap values iterator.
+#[unstable(feature = "map_values_mut", reason = "recently added", issue = "32551")]
+pub struct ValuesMut<'a, K: 'a, V: 'a> {
+    inner: IterMut<'a, K, V>
+}
+
 enum InternalEntry<K, V, M> {
     Occupied {
         elem: FullBucket<K, V, M>,
@@ -1460,6 +1494,18 @@ impl<'a, K, V> ExactSizeIterator for Values<'a, K, V> {
     #[inline] fn len(&self) -> usize { self.inner.len() }
 }
 
+#[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;
+
+    #[inline] fn next(&mut self) -> Option<(&'a mut V)> { self.inner.next().map(|(_, v)| v) }
+    #[inline] 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> ExactSizeIterator for ValuesMut<'a, K, V> {
+    #[inline] fn len(&self) -> usize { self.inner.len() }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a, K, V> Iterator for Drain<'a, K, V> {
     type Item = (K, V);
@@ -1907,6 +1953,7 @@ mod test_map {
         assert_eq!(m.drain().next(), None);
         assert_eq!(m.keys().next(), None);
         assert_eq!(m.values().next(), None);
+        assert_eq!(m.values_mut().next(), None);
         assert_eq!(m.iter().next(), None);
         assert_eq!(m.iter_mut().next(), None);
         assert_eq!(m.len(), 0);
@@ -2084,6 +2131,20 @@ mod test_map {
     }
 
     #[test]
+    fn test_values_mut() {
+        let vec = vec![(1, 1), (2, 2), (3, 3)];
+        let mut map: HashMap<_, _> = vec.into_iter().collect();
+        for value in map.values_mut() {
+            *value = (*value) * 2
+        }
+        let values: Vec<_> = map.values().cloned().collect();
+        assert_eq!(values.len(), 3);
+        assert!(values.contains(&2));
+        assert!(values.contains(&4));
+        assert!(values.contains(&6));
+    }
+
+    #[test]
     fn test_find() {
         let mut m = HashMap::new();
         assert!(m.get(&1).is_none());
diff --git a/src/libstd/lib.rs b/src/libstd/lib.rs
index f7b1042592d..8dcac514172 100644
--- a/src/libstd/lib.rs
+++ b/src/libstd/lib.rs
@@ -241,6 +241,7 @@
 #![feature(link_args)]
 #![feature(linkage)]
 #![feature(macro_reexport)]
+#![cfg_attr(test, feature(map_values_mut))]
 #![feature(num_bits_bytes)]
 #![feature(old_wrapping)]
 #![feature(on_unimplemented)]