1
1
use std:: borrow:: Borrow ;
2
- use std:: collections:: hash_map:: RawEntryMut ;
3
2
use std:: hash:: { Hash , Hasher } ;
4
- use std:: iter;
3
+ use std:: { iter, mem } ;
5
4
6
5
use either:: Either ;
6
+ use hashbrown:: hash_table:: { Entry , HashTable } ;
7
7
8
- use crate :: fx:: { FxHashMap , FxHasher } ;
8
+ use crate :: fx:: FxHasher ;
9
9
use crate :: sync:: { CacheAligned , Lock , LockGuard , Mode , is_dyn_thread_safe} ;
10
10
11
11
// 32 shards is sufficient to reduce contention on an 8-core Ryzen 7 1700,
@@ -140,17 +140,67 @@ pub fn shards() -> usize {
140
140
1
141
141
}
142
142
143
- pub type ShardedHashMap < K , V > = Sharded < FxHashMap < K , V > > ;
143
+ pub type ShardedHashMap < K , V > = Sharded < HashTable < ( K , V ) > > ;
144
144
145
145
impl < K : Eq , V > ShardedHashMap < K , V > {
146
146
pub fn with_capacity ( cap : usize ) -> Self {
147
- Self :: new ( || FxHashMap :: with_capacity_and_hasher ( cap, rustc_hash :: FxBuildHasher :: default ( ) ) )
147
+ Self :: new ( || HashTable :: with_capacity ( cap) )
148
148
}
149
149
pub fn len ( & self ) -> usize {
150
150
self . lock_shards ( ) . map ( |shard| shard. len ( ) ) . sum ( )
151
151
}
152
152
}
153
153
154
+ impl < K : Eq + Hash , V > ShardedHashMap < K , V > {
155
+ #[ inline]
156
+ pub fn get < Q > ( & self , key : & Q ) -> Option < V >
157
+ where
158
+ K : Borrow < Q > ,
159
+ Q : Hash + Eq ,
160
+ V : Clone ,
161
+ {
162
+ let hash = make_hash ( key) ;
163
+ let shard = self . lock_shard_by_hash ( hash) ;
164
+ let ( _, value) = shard. find ( hash, |( k, _) | k. borrow ( ) == key) ?;
165
+ Some ( value. clone ( ) )
166
+ }
167
+
168
+ #[ inline]
169
+ pub fn get_or_insert_with ( & self , key : K , default : impl FnOnce ( ) -> V ) -> V
170
+ where
171
+ V : Copy ,
172
+ {
173
+ let hash = make_hash ( & key) ;
174
+ let mut shard = self . lock_shard_by_hash ( hash) ;
175
+
176
+ match table_entry ( & mut shard, hash, & key) {
177
+ Entry :: Occupied ( e) => e. get ( ) . 1 ,
178
+ Entry :: Vacant ( e) => {
179
+ let value = default ( ) ;
180
+ e. insert ( ( key, value) ) ;
181
+ value
182
+ }
183
+ }
184
+ }
185
+
186
+ #[ inline]
187
+ pub fn insert ( & self , key : K , value : V ) -> Option < V > {
188
+ let hash = make_hash ( & key) ;
189
+ let mut shard = self . lock_shard_by_hash ( hash) ;
190
+
191
+ match table_entry ( & mut shard, hash, & key) {
192
+ Entry :: Occupied ( e) => {
193
+ let previous = mem:: replace ( & mut e. into_mut ( ) . 1 , value) ;
194
+ Some ( previous)
195
+ }
196
+ Entry :: Vacant ( e) => {
197
+ e. insert ( ( key, value) ) ;
198
+ None
199
+ }
200
+ }
201
+ }
202
+ }
203
+
154
204
impl < K : Eq + Hash + Copy > ShardedHashMap < K , ( ) > {
155
205
#[ inline]
156
206
pub fn intern_ref < Q : ?Sized > ( & self , value : & Q , make : impl FnOnce ( ) -> K ) -> K
@@ -160,13 +210,12 @@ impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
160
210
{
161
211
let hash = make_hash ( value) ;
162
212
let mut shard = self . lock_shard_by_hash ( hash) ;
163
- let entry = shard. raw_entry_mut ( ) . from_key_hashed_nocheck ( hash, value) ;
164
213
165
- match entry {
166
- RawEntryMut :: Occupied ( e) => * e . key ( ) ,
167
- RawEntryMut :: Vacant ( e) => {
214
+ match table_entry ( & mut shard , hash , value ) {
215
+ Entry :: Occupied ( e) => e . get ( ) . 0 ,
216
+ Entry :: Vacant ( e) => {
168
217
let v = make ( ) ;
169
- e. insert_hashed_nocheck ( hash , v, ( ) ) ;
218
+ e. insert ( ( v, ( ) ) ) ;
170
219
v
171
220
}
172
221
}
@@ -180,13 +229,12 @@ impl<K: Eq + Hash + Copy> ShardedHashMap<K, ()> {
180
229
{
181
230
let hash = make_hash ( & value) ;
182
231
let mut shard = self . lock_shard_by_hash ( hash) ;
183
- let entry = shard. raw_entry_mut ( ) . from_key_hashed_nocheck ( hash, & value) ;
184
232
185
- match entry {
186
- RawEntryMut :: Occupied ( e) => * e . key ( ) ,
187
- RawEntryMut :: Vacant ( e) => {
233
+ match table_entry ( & mut shard , hash , & value ) {
234
+ Entry :: Occupied ( e) => e . get ( ) . 0 ,
235
+ Entry :: Vacant ( e) => {
188
236
let v = make ( value) ;
189
- e. insert_hashed_nocheck ( hash , v, ( ) ) ;
237
+ e. insert ( ( v, ( ) ) ) ;
190
238
v
191
239
}
192
240
}
@@ -203,17 +251,30 @@ impl<K: Eq + Hash + Copy + IntoPointer> ShardedHashMap<K, ()> {
203
251
let hash = make_hash ( & value) ;
204
252
let shard = self . lock_shard_by_hash ( hash) ;
205
253
let value = value. into_pointer ( ) ;
206
- shard. raw_entry ( ) . from_hash ( hash, |entry| entry . into_pointer ( ) == value) . is_some ( )
254
+ shard. find ( hash, |( k , ( ) ) | k . into_pointer ( ) == value) . is_some ( )
207
255
}
208
256
}
209
257
210
258
#[ inline]
211
- pub fn make_hash < K : Hash + ?Sized > ( val : & K ) -> u64 {
259
+ fn make_hash < K : Hash + ?Sized > ( val : & K ) -> u64 {
212
260
let mut state = FxHasher :: default ( ) ;
213
261
val. hash ( & mut state) ;
214
262
state. finish ( )
215
263
}
216
264
265
+ #[ inline]
266
+ fn table_entry < ' a , K , V , Q > (
267
+ table : & ' a mut HashTable < ( K , V ) > ,
268
+ hash : u64 ,
269
+ key : & Q ,
270
+ ) -> Entry < ' a , ( K , V ) >
271
+ where
272
+ K : Hash + Borrow < Q > ,
273
+ Q : ?Sized + Eq ,
274
+ {
275
+ table. entry ( hash, move |( k, _) | k. borrow ( ) == key, |( k, _) | make_hash ( k) )
276
+ }
277
+
217
278
/// Get a shard with a pre-computed hash value. If `get_shard_by_value` is
218
279
/// ever used in combination with `get_shard_by_hash` on a single `Sharded`
219
280
/// instance, then `hash` must be computed with `FxHasher`. Otherwise,
0 commit comments