diff --git a/src/lib.rs b/src/lib.rs index 6125b4d..d28ec36 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -609,6 +609,58 @@ where self.inner.iter_mut() } + /// An iterator visiting all key-value pairs in arbitrary order. The iterator returns + /// a reference to the key and the first element in the corresponding key's vector. + /// Iterator element type is (&'a K, &'a V). + /// + /// # Examples + /// + /// ``` + /// use multimap::MultiMap; + /// + /// let mut map = MultiMap::new(); + /// map.insert(1,42); + /// map.insert(1,1337); + /// map.insert(3,2332); + /// map.insert(4,1991); + /// + /// for (key, value) in map.flat_iter() { + /// println!("key: {:?}, val: {:?}", key, value); + /// } + /// ``` + pub fn flat_iter(&self) -> FlatIter { + FlatIter { + inner: self.inner.iter(), + inner_iter: None, + } + } + + /// An iterator visiting all key-value pairs in arbitrary order. The iterator returns + /// a reference to the key and the first element in the corresponding key's vector. + /// Iterator element type is (&'a K, &'a V). + /// + /// # Examples + /// + /// ``` + /// use multimap::MultiMap; + /// + /// let mut map = MultiMap::new(); + /// map.insert(1,42); + /// map.insert(1,1337); + /// map.insert(3,2332); + /// map.insert(4,1991); + /// + /// for (key, value) in map.flat_iter_mut() { + /// println!("key: {:?}, val: {:?}", key, value); + /// } + /// ``` + pub fn flat_iter_mut(&mut self) -> FlatIterMut { + FlatIterMut { + inner: self.inner.iter_mut(), + inner_iter: None, + } + } + /// Gets the specified key's corresponding entry in the map for in-place manipulation. /// It's possible to both manipulate the vector and the 'value' (the first value in the /// vector). @@ -894,6 +946,63 @@ impl<'a, K, V> ExactSizeIterator for IterMut<'a, K, V> { } } +#[derive(Clone)] +pub struct FlatIter<'a, K: 'a, V: 'a> { + inner: IterAll<'a, K, Vec>, + inner_iter: Option<(&'a K, std::slice::Iter<'a, V>)>, +} + +impl<'a, K, V> Iterator for FlatIter<'a, K, V> { + type Item = (&'a K, &'a V); + + fn next(&mut self) -> Option { + if let Some((ref k, ref mut iter)) = self.inner_iter { + match iter.next() { + Some(v) => return Some((k, v)), + None => self.inner_iter = None, + } + } + match self.inner.next() { + Some((k, vec)) if vec.len() == 1 => Some((k, &vec[0])), + Some((k, vec)) => { + let mut iter = vec.iter(); + let v = iter.next().unwrap(); + self.inner_iter = Some((k, iter)); + Some((k, v)) + }, + _ => None, + } + } +} + +pub struct FlatIterMut<'a, K: 'a, V: 'a> { + inner: IterAllMut<'a, K, Vec>, + inner_iter: Option<(&'a K, std::slice::IterMut<'a, V>)>, +} + +impl<'a, K, V> Iterator for FlatIterMut<'a, K, V> { + type Item = (&'a K, &'a mut V); + + fn next(&mut self) -> Option { + if let Some((ref k, ref mut iter)) = self.inner_iter { + match iter.next() { + Some(v) => return Some((k, v)), + None => self.inner_iter = None, + } + } + match self.inner.next() { + Some((k, vec)) if vec.len() == 1 => Some((k, &mut vec[0])), + Some((k, vec)) => { + let mut iter = vec.iter_mut(); + let v = iter.next().unwrap(); + self.inner_iter = Some((k, iter)); + Some((k, v)) + }, + _ => None, + } + } +} + #[macro_export] /// Create a `MultiMap` from a list of key value pairs /// @@ -1187,6 +1296,54 @@ mod tests { assert_eq!(iter.len(), 1); } + #[test] + fn flat_iter() { + let mut m: MultiMap = MultiMap::new(); + m.insert(1, 42); + m.insert(1, 43); + m.insert(4, 42); + m.insert(8, 42); + + let keys = vec![1, 4, 8]; + + for (key, value) in m.flat_iter() { + assert!(keys.contains(key)); + + if key == &1 { + assert!(value == &42 || value == &43); + } else { + assert_eq!(value, &42); + } + } + } + + #[test] + fn flat_iter_mut() { + let mut m: MultiMap = MultiMap::new(); + m.insert(1, 42); + m.insert(1, 43); + m.insert(4, 42); + m.insert(8, 42); + + let keys = vec![1, 4, 8]; + + for (key, value) in m.flat_iter_mut() { + assert!(keys.contains(key)); + + if key == &1 { + assert!(value == &42 || value == &43); + + *value = 55; + assert_eq!(value, &55); + } else { + assert_eq!(value, &42); + + *value = 76; + assert_eq!(value, &76); + } + } + } + #[test] fn intoiterator_for_reference_type() { let mut m: MultiMap = MultiMap::new();