diff --git a/src/generic_containers.rs b/src/generic_containers.rs index b217987e5..339f7a9a7 100644 --- a/src/generic_containers.rs +++ b/src/generic_containers.rs @@ -15,6 +15,15 @@ pub trait Map { type Value; fn insert(&mut self, key: Self::Key, value: Self::Value) -> Option; fn remove(&mut self, key: &Self::Key) -> Option; + fn aggregate(&mut self, key: Self::Key, t: T, mut operation: F) + where + F: FnMut(Option, &Self::Key, T) -> Option, + { + let opt_value = self.remove(&key); + if let Some(value) = operation(opt_value, &key, t) { + self.insert(key, value); + } + } fn entry_or_default(&mut self, key: Self::Key) -> &mut Self::Value where Self::Value: Default; @@ -81,6 +90,16 @@ where let index = self.iter().position(|(k, _)| k == key)?; Some(self.swap_remove(index).1) } + fn aggregate(&mut self, key: K, t: T, mut operation: F) + where + F: FnMut(Option, &K, T) -> Option, + { + let opt_value = Map::remove(self, &key); + if let Some(value) = operation(opt_value, &key, t) { + // The key was removed so a single push is enough to insert it back. + self.push((key, value)); + } + } fn entry_or_default(&mut self, key: K) -> &mut V where V: Default, diff --git a/src/grouping_map.rs b/src/grouping_map.rs index 9f0220ea2..0f6314934 100644 --- a/src/grouping_map.rs +++ b/src/grouping_map.rs @@ -144,12 +144,8 @@ where { let mut destination_map = self.map; - self.iter.for_each(|(key, val)| { - let acc = destination_map.remove(&key); - if let Some(op_res) = operation(acc, &key, val) { - destination_map.insert(key, op_res); - } - }); + self.iter + .for_each(|(key, val)| destination_map.aggregate(key, val, &mut operation)); destination_map }