freya_core/
extended_hashmap.rs

1use std::{
2    collections::{
3        HashMap,
4        HashSet,
5        hash_map::Entry,
6    },
7    hash::Hash,
8};
9
10pub trait ExtendedHashMap<K, V> {
11    fn get_disjoint_entries<const N: usize>(
12        &mut self,
13        entries: [&K; N],
14        default: impl FnMut(&K) -> V,
15    ) -> Option<[&mut V; N]>;
16
17    fn get_disjoint_two_entries(
18        &mut self,
19        left: &K,
20        right: &K,
21        left_default: impl FnMut(&K) -> V,
22        right_default: impl FnMut(&V, &K) -> V,
23    ) -> [Option<&mut V>; 2]
24    where
25        V: Clone;
26}
27
28impl<K: Eq + Hash + ToOwned<Owned = K>, V, S: std::hash::BuildHasher + std::default::Default>
29    ExtendedHashMap<K, V> for HashMap<K, V, S>
30{
31    fn get_disjoint_entries<const N: usize>(
32        &mut self,
33        entries: [&K; N],
34        mut default: impl FnMut(&K) -> V,
35    ) -> Option<[&mut V; N]> {
36        let keys = HashSet::<&K, S>::from_iter(entries);
37
38        if keys.len() != N {
39            return None;
40        }
41
42        Some(entries.map(|key| {
43            let ptr: *mut V = match self.entry(key.to_owned()) {
44                Entry::Occupied(e) => e.into_mut(),
45                Entry::Vacant(e) => e.insert(default(key)),
46            };
47            unsafe { &mut *ptr }
48        }))
49    }
50
51    fn get_disjoint_two_entries(
52        &mut self,
53        left: &K,
54        right: &K,
55        mut left_default: impl FnMut(&K) -> V,
56        mut right_default: impl FnMut(&V, &K) -> V,
57    ) -> [Option<&mut V>; 2]
58    where
59        V: Clone,
60    {
61        let left_val = self
62            .entry(left.to_owned())
63            .or_insert_with(|| left_default(left))
64            .clone();
65
66        self.entry(right.to_owned())
67            .or_insert_with(|| right_default(&left_val, right));
68
69        self.get_disjoint_mut([left, right])
70    }
71}
freya_core/extended_hashmap.rs

freya_core/
extended_hashmap.rs
