diff --git a/primitives/core/src/bounded/bounded_btree_map.rs b/primitives/core/src/bounded/bounded_btree_map.rs index 32f8d85ca795b..d2c148d6de9c5 100644 --- a/primitives/core/src/bounded/bounded_btree_map.rs +++ b/primitives/core/src/bounded/bounded_btree_map.rs @@ -168,6 +168,37 @@ where pub fn iter_mut(&mut self) -> sp_std::collections::btree_map::IterMut { self.0.iter_mut() } + + /// Consume the map, applying `f` to each of it's values and returning a new map. + pub fn map(self, mut f: F) -> BoundedBTreeMap + where + F: FnMut((&K, V)) -> T, + { + BoundedBTreeMap::::unchecked_from( + self.0 + .into_iter() + .map(|(k, v)| { + let t = f((&k, v)); + (k, t) + }) + .collect(), + ) + } + + /// Consume the map, applying `f` to each of it's values as long as it returns successfully. If + /// an `Err(E)` is ever encountered, the mapping is short circuited and the error is returned; + /// otherwise, a new map is returned in the contained `Ok` value. + pub fn try_map(self, mut f: F) -> Result, E> + where + F: FnMut((&K, V)) -> Result, + { + Ok(BoundedBTreeMap::::unchecked_from( + self.0 + .into_iter() + .map(|(k, v)| (f((&k, v)).map(|t| (k, t)))) + .collect::, _>>()?, + )) + } } impl Default for BoundedBTreeMap @@ -537,4 +568,58 @@ pub mod test { assert_eq!(b1, b2); } + + #[test] + fn map_retains_size() { + let b1 = boundedmap_from_keys::>(&[1, 2]); + let b2 = b1.clone(); + + assert_eq!(b1.len(), b2.map(|(_, _)| 5_u32).len()); + } + + #[test] + fn map_maps_properly() { + let b1: BoundedBTreeMap> = + [1, 2, 3, 4].into_iter().map(|k| (k, k * 2)).try_collect().unwrap(); + let b2: BoundedBTreeMap> = + [1, 2, 3, 4].into_iter().map(|k| (k, k)).try_collect().unwrap(); + + assert_eq!(b1, b2.map(|(_, v)| v * 2)); + } + + #[test] + fn try_map_retains_size() { + let b1 = boundedmap_from_keys::>(&[1, 2]); + let b2 = b1.clone(); + + assert_eq!(b1.len(), b2.try_map::<_, (), _>(|(_, _)| Ok(5_u32)).unwrap().len()); + } + + #[test] + fn try_map_maps_properly() { + let b1: BoundedBTreeMap> = + [1, 2, 3, 4].into_iter().map(|k| (k, k * 2)).try_collect().unwrap(); + let b2: BoundedBTreeMap> = + [1, 2, 3, 4].into_iter().map(|k| (k, k)).try_collect().unwrap(); + + assert_eq!(b1, b2.try_map::<_, (), _>(|(_, v)| Ok(v * 2)).unwrap()); + } + + #[test] + fn try_map_short_circuit() { + let b1: BoundedBTreeMap> = + [1, 2, 3, 4].into_iter().map(|k| (k, k)).try_collect().unwrap(); + + assert_eq!(Err("overflow"), b1.try_map(|(_, v)| v.checked_mul(100).ok_or("overflow"))); + } + + #[test] + fn try_map_ok() { + let b1: BoundedBTreeMap> = + [1, 2, 3, 4].into_iter().map(|k| (k, k)).try_collect().unwrap(); + let b2: BoundedBTreeMap> = + [1, 2, 3, 4].into_iter().map(|k| (k, (k as u16) * 100)).try_collect().unwrap(); + + assert_eq!(Ok(b2), b1.try_map(|(_, v)| (v as u16).checked_mul(100_u16).ok_or("overflow"))); + } }