Implement ManagedMap::iter_mut().

This commit is contained in:
Astro 2018-03-24 04:21:46 +01:00 committed by whitequark
parent ecbc2ee8e5
commit 0192eb98d8
2 changed files with 79 additions and 5 deletions

View File

@ -19,4 +19,5 @@ pub use object::Managed;
pub use slice::ManagedSlice; pub use slice::ManagedSlice;
#[cfg(feature = "map")] #[cfg(feature = "map")]
pub use map::{ManagedMap, pub use map::{ManagedMap,
Iter as ManagedMapIter}; Iter as ManagedMapIter,
IterMut as ManagedMapIterMut};

View File

@ -6,11 +6,11 @@ use core::borrow::Borrow;
#[cfg(feature = "std")] #[cfg(feature = "std")]
use std::collections::BTreeMap; use std::collections::BTreeMap;
#[cfg(feature = "std")] #[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")))] #[cfg(all(feature = "alloc", not(feature = "std")))]
use alloc::btree_map::BTreeMap; use alloc::btree_map::BTreeMap;
#[cfg(all(feature = "alloc", not(feature = "std")))] #[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. /// 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 { match self {
&ManagedMap::Borrowed(ref pairs) => &ManagedMap::Borrowed(ref pairs) =>
Iter::Borrowed(pairs.iter()), Iter::Borrowed(pairs.iter()),
@ -235,6 +235,16 @@ impl<'a, K: Ord + 'a, V: 'a> ManagedMap<'a, K, V> {
Iter::Owned(map.iter()), 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> { 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)>>), Borrowed(slice::Iter<'a, Option<(K, V)>>),
/// Owned variant, only available with the `std` or `alloc` feature enabled. /// Owned variant, only available with the `std` or `alloc` feature enabled.
#[cfg(any(feature = "std", feature = "alloc"))] #[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> { 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 // LCOV_EXCL_START
#[cfg(test)] #[cfg(test)]
mod test { mod test {
@ -456,4 +504,29 @@ mod test {
assert_eq!(iter.next(), Some((&"d", &4))); assert_eq!(iter.next(), Some((&"d", &4)));
assert_eq!(iter.next(), None); 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);
}
}
} }