Skip to main content

stefans_utils/
collect_by_key.rs

1pub trait CollectByKey<Item, Key, Output> {
2    fn collect_by_key(self, f: impl Fn(&Item) -> Key) -> Output;
3}
4
5macro_rules! impl_for {
6    ($map_type:ty, $list_type:ty, $add_to_list_method:ident) => {
7        impl_for!($map_type, $list_type, $add_to_list_method, From::<I>);
8    };
9
10    ($map_type:ty, $list_type:ty, $add_to_list_method:ident, $item_type_bounds:expr) => {
11        paste::paste! {
12            impl<T: IntoIterator<Item = I>, I: $item_type_bounds, K: Eq + std::hash::Hash> CollectByKey<I, K, $map_type<K, $list_type<I>>> for T {
13                fn collect_by_key(self, f: impl Fn(&I) -> K) -> $map_type<K, $list_type<I>> {
14                    self.into_iter()
15                        .fold($map_type::default(), |mut acc, curr| {
16                            let key = f(&curr);
17                            match acc.get_mut(&key) {
18                                Some(list) => {
19                                    list.$add_to_list_method(curr);
20                                }
21                                None => {
22                                    let mut list = $list_type::default();
23                                    list.$add_to_list_method(curr);
24                                    acc.insert(key, list);
25                                }
26                            }
27                            acc
28                        })
29                }
30            }
31        }
32    };
33}
34
35impl_for!(std::collections::HashMap, Vec, push);
36impl_for!(
37    std::collections::HashMap,
38    std::collections::HashSet,
39    insert,
40    Eq + std::hash::Hash
41);
42#[cfg(feature = "indexmap")]
43impl_for!(
44    std::collections::HashMap,
45    indexmap::IndexSet,
46    insert,
47    Eq + std::hash::Hash
48);
49
50#[cfg(feature = "indexmap")]
51impl_for!(indexmap::IndexMap, Vec, push);
52#[cfg(feature = "indexmap")]
53impl_for!(
54    indexmap::IndexMap,
55    std::collections::HashSet,
56    insert,
57    Eq + std::hash::Hash
58);
59#[cfg(feature = "indexmap")]
60impl_for!(
61    indexmap::IndexMap,
62    indexmap::IndexSet,
63    insert,
64    Eq + std::hash::Hash
65);