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 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 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 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 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}