Skip to main content

graft_core/
hash_table.rs

1use foldhash::fast::RandomState;
2use hashbrown::hash_table::{AbsentEntry, Entry, IntoIter, Iter, IterMut, OccupiedEntry};
3use std::{
4    borrow::Borrow,
5    hash::{BuildHasher, Hash},
6};
7
8pub trait HTEntry {
9    type Key: Hash + Clone + PartialEq + Eq;
10    fn key(&self) -> &Self::Key;
11}
12
13pub trait HTFromKey: HTEntry {
14    fn from_key(key: Self::Key) -> Self;
15}
16
17#[derive(Debug, Clone)]
18pub struct HashTable<T: HTEntry> {
19    table: hashbrown::HashTable<T>,
20    hash_builder: RandomState,
21}
22
23impl<T: HTEntry> Default for HashTable<T> {
24    fn default() -> Self {
25        Self {
26            table: Default::default(),
27            hash_builder: RandomState::default(),
28        }
29    }
30}
31
32impl<T: HTFromKey> HashTable<T> {
33    /// Returns a mutable reference to the value corresponding to the key.
34    /// If the key is not present in the table, it is inserted with a default value.
35    pub fn get_mut(&mut self, key: &T::Key) -> &mut T {
36        self.entry(key)
37            .or_insert_with(|| T::from_key(key.clone()))
38            .into_mut()
39    }
40
41    /// Ensures a value for the key by inserting the default value if it is not present.
42    pub fn ensure(&mut self, key: &T::Key) {
43        self.entry(key).or_insert_with(|| T::from_key(key.clone()));
44    }
45}
46
47impl<T: HTEntry> HashTable<T> {
48    fn h<K>(&self, key: &K) -> u64
49    where
50        T::Key: Borrow<K> + PartialEq<K>,
51        K: Hash + Eq + ?Sized,
52    {
53        self.hash_builder.hash_one(key)
54    }
55
56    pub fn len(&self) -> usize {
57        self.table.len()
58    }
59
60    pub fn is_empty(&self) -> bool {
61        self.table.is_empty()
62    }
63
64    pub fn has<K>(&self, key: &K) -> bool
65    where
66        T::Key: Borrow<K> + PartialEq<K>,
67        K: Hash + Eq + ?Sized,
68    {
69        self.find(key).is_some()
70    }
71
72    /// Finds a reference to the value corresponding to the key
73    /// or returns `None` if the key is not present in the table.
74    pub fn find<K>(&self, key: &K) -> Option<&T>
75    where
76        T::Key: Borrow<K> + PartialEq<K>,
77        K: Hash + Eq + ?Sized,
78    {
79        self.table.find(self.h(key), |entry| entry.key() == key)
80    }
81
82    /// Finds a mutable reference to the value corresponding to the key
83    /// or returns `None` if the key is not present in the table.
84    pub fn find_mut<K>(&mut self, key: &K) -> Option<&mut T>
85    where
86        T::Key: Borrow<K> + PartialEq<K>,
87        K: Hash + Eq + ?Sized,
88    {
89        self.table.find_mut(self.h(key), |entry| entry.key() == key)
90    }
91
92    pub fn insert(&mut self, item: T) {
93        match self.entry(item.key()) {
94            Entry::Occupied(mut entry) => {
95                *entry.get_mut() = item;
96            }
97            Entry::Vacant(entry) => {
98                entry.insert(item);
99            }
100        }
101    }
102
103    pub fn extract_if<'a, F>(&'a mut self, f: F) -> impl Iterator<Item = T> + 'a
104    where
105        F: FnMut(&mut T) -> bool + 'a,
106    {
107        self.table.extract_if(f)
108    }
109
110    pub fn remove<K>(&mut self, key: &K) -> Option<T>
111    where
112        T::Key: Borrow<K> + PartialEq<K>,
113        K: Hash + Eq + ?Sized,
114    {
115        if let Ok(entry) = self
116            .table
117            .find_entry(self.h(key), |entry| entry.key() == key)
118        {
119            let (v, _) = entry.remove();
120            Some(v)
121        } else {
122            None
123        }
124    }
125
126    pub fn find_entry<K>(&mut self, key: &K) -> Result<OccupiedEntry<'_, T>, AbsentEntry<'_, T>>
127    where
128        T::Key: Borrow<K> + PartialEq<K>,
129        K: Hash + Eq + ?Sized,
130    {
131        self.table
132            .find_entry(self.hash_builder.hash_one(key), |entry| entry.key() == key)
133    }
134
135    pub fn entry<K>(&mut self, key: &K) -> Entry<'_, T>
136    where
137        T::Key: Borrow<K> + PartialEq<K>,
138        K: Hash + Eq + ?Sized,
139    {
140        let hb = &self.hash_builder;
141        self.table.entry(
142            hb.hash_one(key),
143            |entry| entry.key() == key,
144            |entry| hb.hash_one(entry.key()),
145        )
146    }
147
148    #[inline]
149    pub fn iter(&self) -> Iter<'_, T> {
150        self.table.iter()
151    }
152
153    #[inline]
154    pub fn iter_mut(&mut self) -> IterMut<'_, T> {
155        self.table.iter_mut()
156    }
157
158    pub fn first_key(&self) -> Option<T::Key> {
159        self.iter().next().map(|entry| entry.key().clone())
160    }
161}
162
163impl<T: HTEntry> IntoIterator for HashTable<T> {
164    type Item = T;
165    type IntoIter = IntoIter<Self::Item>;
166
167    fn into_iter(self) -> Self::IntoIter {
168        self.table.into_iter()
169    }
170}
171
172impl<'a, T: HTEntry> IntoIterator for &'a HashTable<T> {
173    type Item = &'a T;
174    type IntoIter = Iter<'a, T>;
175
176    #[inline]
177    fn into_iter(self) -> Self::IntoIter {
178        self.iter()
179    }
180}
181
182impl<'a, T: HTEntry> IntoIterator for &'a mut HashTable<T> {
183    type Item = &'a mut T;
184    type IntoIter = IterMut<'a, T>;
185
186    #[inline]
187    fn into_iter(self) -> Self::IntoIter {
188        self.iter_mut()
189    }
190}