utils.rs

pub trait IteratorExt: std::iter::Iterator {
    fn group_by<F, K, G>(self, f: F) -> HashMap<K, Vec<G>>
    where
        Self: Sized,
        K: Eq + Hash,
        F: FnMut(Self::Item) -> (K, G);
}

impl<T> IteratorExt for T
where
    T: std::iter::Iterator + ?Sized,
{
    fn group_by<F, K, G>(self, mut f: F) -> HashMap<K, Vec<G>>
    where
        Self: Sized,
        K: Eq + Hash,
        F: FnMut(Self::Item) -> (K, G),
    {
        let mut accum = HashMap::<K, Vec<G>>::new();
        for item in self {
            let (k, g) = f(item);
            if let Some(group) = accum.get_mut(&k) {
                group.push(g);
            } else {
                accum.insert(k, vec![g]);
            }
        }
        accum
    }
}