Skip to main content

kaspa_utils/
hashmap.rs

1use std::collections::{HashMap, hash_map::Entry};
2use std::hash::Hash;
3use std::vec::Vec;
4
5pub trait NestedHashMapExtensions<K, K2, V> {
6    /// For use with a hashmap within a hashmap, takes an outer and inner key, and an inner value
7    /// 1) if the outer key exists with and inner hashmap the key-value pair is inserted to the existing hashmap.
8    /// 2) a new inner hashmap is created with the key-value pair if the outer key does not exist.
9    fn insert_into_nested(&mut self, outer_key: K, inner_key: K2, inner_value: V);
10}
11
12impl<K, K2, V> NestedHashMapExtensions<K, K2, V> for HashMap<K, HashMap<K2, V>>
13where
14    K: Hash + PartialEq + Eq,
15    K2: Hash + PartialEq + Eq,
16{
17    fn insert_into_nested(&mut self, outer_key: K, inner_key: K2, inner_value: V) {
18        match self.entry(outer_key) {
19            Entry::Occupied(mut entry) => {
20                entry.get_mut().insert(inner_key, inner_value);
21            }
22            Entry::Vacant(entry) => {
23                let mut inner = HashMap::with_capacity(1);
24                // We do not expect duplicate insert since we just created the hashmap.
25                let _ = &inner.insert(inner_key, inner_value);
26                entry.insert(inner);
27            }
28        }
29    }
30}
31
32pub trait GroupExtension<K, V, I>
33where
34    K: std::hash::Hash + Ord,
35    I: IntoIterator<Item = (K, V)>,
36{
37    fn group_from(v: I) -> HashMap<K, Vec<V>>;
38}
39
40impl<K, V, I> GroupExtension<K, V, I> for HashMap<K, Vec<V>>
41where
42    K: std::hash::Hash + Ord,
43    I: IntoIterator<Item = (K, V)>,
44{
45    fn group_from(v: I) -> HashMap<K, Vec<V>> {
46        let mut result = HashMap::<K, Vec<V>>::new();
47        for (a, b) in v {
48            result.entry(a).or_default().push(b);
49        }
50        result
51    }
52}