gpgrv/
hash_multimap.rs

1use std::collections::HashMap;
2use std::collections::HashSet;
3use std::hash::Hash;
4
5pub struct HashMultiMap<K, V> {
6    inner: HashMap<K, HashSet<V>>,
7}
8
9impl<K: Eq + Hash, V: Eq + Hash> HashMultiMap<K, V> {
10    pub fn new() -> Self {
11        HashMultiMap {
12            inner: HashMap::new(),
13        }
14    }
15
16    pub fn insert(&mut self, k: K, v: V) -> bool {
17        self.inner.entry(k).or_insert_with(HashSet::new).insert(v)
18    }
19
20    pub fn values(&self) -> HashSet<&V> {
21        let mut ret = HashSet::with_capacity(self.inner.len());
22        for value_set in self.inner.values() {
23            for value in value_set {
24                ret.insert(value);
25            }
26        }
27        ret
28    }
29
30    pub fn entries<'h>(&'h self) -> Box<dyn Iterator<Item = (&'h K, &'h V)> + 'h> {
31        Box::new(
32            self.inner
33                .iter()
34                .flat_map(|(k, v)| v.iter().map(move |v| (k, v))),
35        )
36    }
37}
38
39impl<K: Clone, V: Clone> Clone for HashMultiMap<K, V> {
40    fn clone(&self) -> Self {
41        HashMultiMap {
42            inner: self.inner.clone(),
43        }
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use std::collections::HashSet;
50    #[test]
51    fn smoke() {
52        let mut map = super::HashMultiMap::new();
53        assert!(map.insert(5, 2));
54        assert!(map.insert(6, 7));
55        assert!(map.insert(6, 8));
56        assert!(!map.insert(6, 8));
57
58        assert_eq!(3, map.entries().count());
59        assert_eq!(
60            vec![(5, 2), (6, 7), (6, 8)]
61                .into_iter()
62                .collect::<HashSet<(u64, u64)>>(),
63            map.entries().map(|(k, v)| (*k, *v)).collect()
64        );
65    }
66}