Skip to content

Commit

Permalink
Map::aggregate
Browse files Browse the repository at this point in the history
rust-itertools#901 (comment)
SkiFire13 had the idea this could be optimized for vectors.
  • Loading branch information
Philippe-Cholet committed Mar 21, 2024
1 parent dc15e5c commit 2203cb6
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 6 deletions.
19 changes: 19 additions & 0 deletions src/generic_containers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,15 @@ pub trait Map {
type Value;
fn insert(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
fn remove(&mut self, key: &Self::Key) -> Option<Self::Value>;
fn aggregate<T, F>(&mut self, key: Self::Key, t: T, mut operation: F)
where
F: FnMut(Option<Self::Value>, &Self::Key, T) -> Option<Self::Value>,
{
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;
Expand Down Expand Up @@ -81,6 +90,16 @@ where
let index = self.iter().position(|(k, _)| k == key)?;
Some(self.swap_remove(index).1)
}
fn aggregate<T, F>(&mut self, key: K, t: T, mut operation: F)
where
F: FnMut(Option<V>, &K, T) -> Option<V>,
{
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,
Expand Down
8 changes: 2 additions & 6 deletions src/grouping_map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down

0 comments on commit 2203cb6

Please sign in to comment.