earley/
group_by.rs

1use std::collections::BTreeMap;
2use std::collections::btree_map::IntoIter as BTreeIter;
3
4pub trait GroupByExt<I> {
5    fn group_by<K, F>(self, f: F) -> BTreeIter<K, Vec<I>> where
6        F: for<'a> Fn(&'a I) -> K,
7        K: Ord;
8}
9
10impl<T, I> GroupByExt<I> for T where T: Iterator<Item=I> {
11    fn group_by<K, F>(self, f: F) -> BTreeIter<K, Vec<I>> where
12        F: for<'a> Fn(&'a I) -> K,
13        K: Ord
14    {
15        let mut map = BTreeMap::new();
16        for item in self {
17            let hash = f(&item);
18            map.entry(hash).or_insert(vec![]).push(item);
19        }
20        map.into_iter()
21    }
22}
23
24#[test]
25fn test_group_by() {
26    let map: BTreeMap<u32, Vec<u32>> = (1..7).group_by(|i| i%3).collect();
27
28    assert_eq!(map[&0], vec![3, 6]);
29    assert_eq!(map[&1], vec![1, 4]);
30    assert_eq!(map[&2], vec![2, 5]);
31}