From 0192eb98d8899924e3232cb04d769ad641fc0ece Mon Sep 17 00:00:00 2001
From: Astro <astro@spaceboyz.net>
Date: Sat, 24 Mar 2018 04:21:46 +0100
Subject: [PATCH] Implement ManagedMap::iter_mut().

---
 src/lib.rs |  3 +-
 src/map.rs | 81 +++++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0b2631c..b1a71a9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -19,4 +19,5 @@ pub use object::Managed;
 pub use slice::ManagedSlice;
 #[cfg(feature = "map")]
 pub use map::{ManagedMap,
-              Iter as ManagedMapIter};
+              Iter as ManagedMapIter,
+              IterMut as ManagedMapIterMut};
diff --git a/src/map.rs b/src/map.rs
index de7eaea..a9614a7 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -6,11 +6,11 @@ use core::borrow::Borrow;
 #[cfg(feature = "std")]
 use std::collections::BTreeMap;
 #[cfg(feature = "std")]
-use std::collections::btree_map::Iter as BTreeIter;
+use std::collections::btree_map::{Iter as BTreeIter, IterMut as BTreeIterMut};
 #[cfg(all(feature = "alloc", not(feature = "std")))]
 use alloc::btree_map::BTreeMap;
 #[cfg(all(feature = "alloc", not(feature = "std")))]
-use alloc::btree_map::Iter as BTreeIter;
+use alloc::btree_map::{Iter as BTreeIter, IterMut as BTreeIterMut};
 
 /// A managed map.
 ///
@@ -226,7 +226,7 @@ impl<'a, K: Ord + 'a, V: 'a> ManagedMap<'a, K, V> {
         }
     }
 
-    pub fn iter(&'a self) -> Iter<'a, K, V> {
+    pub fn iter(&self) -> Iter<K, V> {
         match self {
             &ManagedMap::Borrowed(ref pairs) =>
                 Iter::Borrowed(pairs.iter()),
@@ -235,6 +235,16 @@ impl<'a, K: Ord + 'a, V: 'a> ManagedMap<'a, K, V> {
                 Iter::Owned(map.iter()),
         }
     }
+
+    pub fn iter_mut(&mut self) -> IterMut<K, V> {
+        match self {
+            &mut ManagedMap::Borrowed(ref mut pairs) =>
+                IterMut::Borrowed(pairs.iter_mut()),
+            #[cfg(any(feature = "std", feature = "alloc"))]
+            &mut ManagedMap::Owned(ref mut map) =>
+                IterMut::Owned(map.iter_mut()),
+        }
+    }
 }
 
 pub enum Iter<'a, K: 'a, V: 'a> {
@@ -242,7 +252,7 @@ pub enum Iter<'a, K: 'a, V: 'a> {
     Borrowed(slice::Iter<'a, Option<(K, V)>>),
     /// Owned variant, only available with the `std` or `alloc` feature enabled.
     #[cfg(any(feature = "std", feature = "alloc"))]
-    Owned(BTreeIter<'a, K, V>)
+    Owned(BTreeIter<'a, K, V>),
 }
 
 impl<'a, K: Ord + 'a, V: 'a> Iterator for Iter<'a, K, V> {
@@ -277,6 +287,44 @@ impl<'a, K: Ord + 'a, V: 'a> Iterator for Iter<'a, K, V> {
     }
 }
 
+pub enum IterMut<'a, K: 'a, V: 'a> {
+    /// Borrowed variant.
+    Borrowed(slice::IterMut<'a, Option<(K, V)>>),
+    /// Owned variant, only available with the `std` or `alloc` feature enabled.
+    #[cfg(any(feature = "std", feature = "alloc"))]
+    Owned(BTreeIterMut<'a, K, V>),
+}
+
+impl<'a, K: Ord + 'a, V: 'a> Iterator for IterMut<'a, K, V> {
+    type Item = (&'a K, &'a mut V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            &mut IterMut::Borrowed(ref mut iter) =>
+                match iter.next() {
+                    Some(&mut Some((ref k, ref mut v))) => Some((&k, v)),
+                    Some(&mut None) => None,
+                    None => None,
+                },
+            #[cfg(any(feature = "std", feature = "alloc"))]
+            &mut IterMut::Owned(ref mut iter) =>
+                iter.next(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        match self {
+            &IterMut::Borrowed(ref iter) => {
+                let (_, upper) = iter.size_hint();
+                (0, upper)
+            },
+            #[cfg(any(feature = "std", feature = "alloc"))]
+            &IterMut::Owned(ref iter) =>
+                iter.size_hint(),
+        }
+    }
+}
+
 // LCOV_EXCL_START
 #[cfg(test)]
 mod test {
@@ -456,4 +504,29 @@ mod test {
         assert_eq!(iter.next(), Some((&"d", &4)));
         assert_eq!(iter.next(), None);
     }
+
+    #[test]
+    fn test_iter_mut_full() {
+        let mut pairs = all_pairs_full();
+        let mut map = ManagedMap::Borrowed(&mut pairs);
+
+        {
+            let mut iter = map.iter_mut();
+            assert_eq!(iter.size_hint(), (0, Some(4)));
+            for (_k, mut v) in &mut iter {
+                *v += 1;
+            }
+            assert_eq!(iter.size_hint(), (0, Some(0)));
+            // Scope for `iter` ends here so that it can be borrowed
+            // again with the following `iter`.
+        }
+        {
+            let mut iter = map.iter();
+            assert_eq!(iter.next(), Some((&"a", &2)));
+            assert_eq!(iter.next(), Some((&"b", &3)));
+            assert_eq!(iter.next(), Some((&"c", &4)));
+            assert_eq!(iter.next(), Some((&"d", &5)));
+            assert_eq!(iter.next(), None);
+        }
+    }
 }