Skip to content

Commit c33ce2d

Browse files
committed
Add the ability to recover the original HashTable from an entry
1 parent 142e536 commit c33ce2d

File tree

2 files changed

+63
-47
lines changed

2 files changed

+63
-47
lines changed

src/external_trait_impls/rayon/table.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ impl<T: Send, A: Allocator> HashTable<T, A> {
165165
#[cfg_attr(feature = "inline-more", inline)]
166166
pub fn par_drain(&mut self) -> ParDrain<'_, T, A> {
167167
ParDrain {
168-
inner: self.table.par_drain(),
168+
inner: self.raw.par_drain(),
169169
}
170170
}
171171
}
@@ -177,7 +177,7 @@ impl<T: Send, A: Allocator + Send> IntoParallelIterator for HashTable<T, A> {
177177
#[cfg_attr(feature = "inline-more", inline)]
178178
fn into_par_iter(self) -> Self::Iter {
179179
IntoParIter {
180-
inner: self.table.into_par_iter(),
180+
inner: self.raw.into_par_iter(),
181181
}
182182
}
183183
}
@@ -189,7 +189,7 @@ impl<'a, T: Sync, A: Allocator> IntoParallelIterator for &'a HashTable<T, A> {
189189
#[cfg_attr(feature = "inline-more", inline)]
190190
fn into_par_iter(self) -> Self::Iter {
191191
ParIter {
192-
inner: unsafe { self.table.par_iter() },
192+
inner: unsafe { self.raw.par_iter() },
193193
marker: PhantomData,
194194
}
195195
}
@@ -202,7 +202,7 @@ impl<'a, T: Send, A: Allocator> IntoParallelIterator for &'a mut HashTable<T, A>
202202
#[cfg_attr(feature = "inline-more", inline)]
203203
fn into_par_iter(self) -> Self::Iter {
204204
ParIterMut {
205-
inner: unsafe { self.table.par_iter() },
205+
inner: unsafe { self.raw.par_iter() },
206206
marker: PhantomData,
207207
}
208208
}

src/table.rs

Lines changed: 59 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ pub struct HashTable<T, A = Global>
4646
where
4747
A: Allocator,
4848
{
49-
pub(crate) table: RawTable<T, A>,
49+
pub(crate) raw: RawTable<T, A>,
5050
}
5151

5252
impl<T> HashTable<T, Global> {
@@ -65,7 +65,7 @@ impl<T> HashTable<T, Global> {
6565
/// ```
6666
pub const fn new() -> Self {
6767
Self {
68-
table: RawTable::new(),
68+
raw: RawTable::new(),
6969
}
7070
}
7171

@@ -84,7 +84,7 @@ impl<T> HashTable<T, Global> {
8484
/// ```
8585
pub fn with_capacity(capacity: usize) -> Self {
8686
Self {
87-
table: RawTable::with_capacity(capacity),
87+
raw: RawTable::with_capacity(capacity),
8888
}
8989
}
9090
}
@@ -133,7 +133,7 @@ where
133133
/// ```
134134
pub const fn new_in(alloc: A) -> Self {
135135
Self {
136-
table: RawTable::new_in(alloc),
136+
raw: RawTable::new_in(alloc),
137137
}
138138
}
139139

@@ -182,13 +182,13 @@ where
182182
/// ```
183183
pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
184184
Self {
185-
table: RawTable::with_capacity_in(capacity, alloc),
185+
raw: RawTable::with_capacity_in(capacity, alloc),
186186
}
187187
}
188188

189189
/// Returns a reference to the underlying allocator.
190190
pub fn allocator(&self) -> &A {
191-
self.table.allocator()
191+
self.raw.allocator()
192192
}
193193

194194
/// Returns a reference to an entry in the table with the given hash and
@@ -222,7 +222,7 @@ where
222222
/// # }
223223
/// ```
224224
pub fn find(&self, hash: u64, eq: impl FnMut(&T) -> bool) -> Option<&T> {
225-
self.table
225+
self.raw
226226
.find(hash, eq)
227227
.map(|bucket| unsafe { bucket.as_ref() })
228228
}
@@ -263,7 +263,7 @@ where
263263
/// # }
264264
/// ```
265265
pub fn find_mut(&mut self, hash: u64, eq: impl FnMut(&T) -> bool) -> Option<&mut T> {
266-
self.table
266+
self.raw
267267
.find(hash, eq)
268268
.map(|bucket| unsafe { bucket.as_mut() })
269269
}
@@ -292,7 +292,7 @@ where
292292
/// let hasher = BuildHasherDefault::<AHasher>::default();
293293
/// let hasher = |val: &_| hasher.hash_one(val);
294294
/// table.insert_unchecked(hasher(&1), (1, "a"), |val| hasher(&val.0));
295-
/// if let Some(entry) = table.find_entry(hasher(&1), |val| val.0 == 1) {
295+
/// if let Ok(entry) = table.find_entry(hasher(&1), |val| val.0 == 1) {
296296
/// entry.remove();
297297
/// }
298298
/// assert_eq!(table.find(hasher(&1), |val| val.0 == 1), None);
@@ -306,12 +306,15 @@ where
306306
&mut self,
307307
hash: u64,
308308
eq: impl FnMut(&T) -> bool,
309-
) -> Option<OccupiedEntry<'_, T, A>> {
310-
self.table.find(hash, eq).map(|bucket| OccupiedEntry {
311-
hash,
312-
bucket,
313-
table: &mut self.table,
314-
})
309+
) -> Result<OccupiedEntry<'_, T, A>, &mut Self> {
310+
match self.raw.find(hash, eq) {
311+
Some(bucket) => Ok(OccupiedEntry {
312+
hash,
313+
bucket,
314+
table: self,
315+
}),
316+
None => Err(self),
317+
}
315318
}
316319

317320
/// Returns an `Entry` for an entry in the table with the given hash
@@ -365,16 +368,16 @@ where
365368
eq: impl FnMut(&T) -> bool,
366369
hasher: impl Fn(&T) -> u64,
367370
) -> Entry<'_, T, A> {
368-
match self.table.find_or_find_insert_slot(hash, eq, hasher) {
371+
match self.raw.find_or_find_insert_slot(hash, eq, hasher) {
369372
Ok(bucket) => Entry::Occupied(OccupiedEntry {
370373
hash,
371374
bucket,
372-
table: &mut self.table,
375+
table: self,
373376
}),
374377
Err(insert_slot) => Entry::Vacant(VacantEntry {
375378
hash,
376379
insert_slot,
377-
table: &mut self.table,
380+
table: self,
378381
}),
379382
}
380383
}
@@ -393,11 +396,11 @@ where
393396
value: T,
394397
hasher: impl Fn(&T) -> u64,
395398
) -> OccupiedEntry<'_, T, A> {
396-
let bucket = self.table.insert(hash, value, hasher);
399+
let bucket = self.raw.insert(hash, value, hasher);
397400
OccupiedEntry {
398401
hash,
399402
bucket,
400-
table: &mut self.table,
403+
table: self,
401404
}
402405
}
403406

@@ -425,7 +428,7 @@ where
425428
/// # }
426429
/// ```
427430
pub fn clear(&mut self) {
428-
self.table.clear();
431+
self.raw.clear();
429432
}
430433

431434
/// Shrinks the capacity of the table as much as possible. It will drop
@@ -459,7 +462,7 @@ where
459462
/// # }
460463
/// ```
461464
pub fn shrink_to_fit(&mut self, hasher: impl Fn(&T) -> u64) {
462-
self.table.shrink_to(self.len(), hasher)
465+
self.raw.shrink_to(self.len(), hasher)
463466
}
464467

465468
/// Shrinks the capacity of the table with a lower limit. It will drop
@@ -498,7 +501,7 @@ where
498501
/// # }
499502
/// ```
500503
pub fn shrink_to(&mut self, min_capacity: usize, hasher: impl Fn(&T) -> u64) {
501-
self.table.shrink_to(min_capacity, hasher);
504+
self.raw.shrink_to(min_capacity, hasher);
502505
}
503506

504507
/// Reserves capacity for at least `additional` more elements to be inserted
@@ -538,7 +541,7 @@ where
538541
/// # }
539542
/// ```
540543
pub fn reserve(&mut self, additional: usize, hasher: impl Fn(&T) -> u64) {
541-
self.table.reserve(additional, hasher)
544+
self.raw.reserve(additional, hasher)
542545
}
543546

544547
/// Tries to reserve capacity for at least `additional` more elements to be inserted
@@ -579,7 +582,7 @@ where
579582
additional: usize,
580583
hasher: impl Fn(&T) -> u64,
581584
) -> Result<(), TryReserveError> {
582-
self.table.try_reserve(additional, hasher)
585+
self.raw.try_reserve(additional, hasher)
583586
}
584587

585588
/// Returns the number of elements the table can hold without reallocating.
@@ -592,7 +595,7 @@ where
592595
/// assert!(table.capacity() >= 100);
593596
/// ```
594597
pub fn capacity(&self) -> usize {
595-
self.table.capacity()
598+
self.raw.capacity()
596599
}
597600

598601
/// Returns the number of elements in the table.
@@ -619,7 +622,7 @@ where
619622
/// # }
620623
/// ```
621624
pub fn len(&self) -> usize {
622-
self.table.len()
625+
self.raw.len()
623626
}
624627

625628
/// Returns `true` if the set contains no elements.
@@ -646,7 +649,7 @@ where
646649
/// # }
647650
/// ```
648651
pub fn is_empty(&self) -> bool {
649-
self.table.is_empty()
652+
self.raw.is_empty()
650653
}
651654

652655
/// An iterator visiting all elements in arbitrary order.
@@ -679,7 +682,7 @@ where
679682
/// ```
680683
pub fn iter(&self) -> Iter<'_, T> {
681684
Iter {
682-
inner: unsafe { self.table.iter() },
685+
inner: unsafe { self.raw.iter() },
683686
marker: PhantomData,
684687
}
685688
}
@@ -731,7 +734,7 @@ where
731734
/// ```
732735
pub fn iter_mut(&mut self) -> IterMut<'_, T> {
733736
IterMut {
734-
inner: unsafe { self.table.iter() },
737+
inner: unsafe { self.raw.iter() },
735738
marker: PhantomData,
736739
}
737740
}
@@ -766,9 +769,9 @@ where
766769
pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
767770
// Here we only use `iter` as a temporary, preventing use-after-free
768771
unsafe {
769-
for item in self.table.iter() {
772+
for item in self.raw.iter() {
770773
if !f(item.as_mut()) {
771-
self.table.erase(item);
774+
self.raw.erase(item);
772775
}
773776
}
774777
}
@@ -807,7 +810,7 @@ where
807810
/// ```
808811
pub fn drain(&mut self) -> Drain<'_, T, A> {
809812
Drain {
810-
inner: self.table.drain(),
813+
inner: self.raw.drain(),
811814
}
812815
}
813816

@@ -858,8 +861,8 @@ where
858861
ExtractIf {
859862
f,
860863
inner: RawExtractIf {
861-
iter: unsafe { self.table.iter() },
862-
table: &mut self.table,
864+
iter: unsafe { self.raw.iter() },
865+
table: &mut self.raw,
863866
},
864867
}
865868
}
@@ -922,7 +925,7 @@ where
922925
hashes: [u64; N],
923926
eq: impl FnMut(usize, &T) -> bool,
924927
) -> Option<[&'_ mut T; N]> {
925-
self.table.get_many_mut(hashes, eq)
928+
self.raw.get_many_mut(hashes, eq)
926929
}
927930

928931
/// Attempts to get mutable references to `N` values in the map at once, without validating that
@@ -992,7 +995,7 @@ where
992995
hashes: [u64; N],
993996
eq: impl FnMut(usize, &T) -> bool,
994997
) -> Option<[&'_ mut T; N]> {
995-
self.table.get_many_unchecked_mut(hashes, eq)
998+
self.raw.get_many_unchecked_mut(hashes, eq)
996999
}
9971000
}
9981001

@@ -1005,7 +1008,7 @@ where
10051008

10061009
fn into_iter(self) -> IntoIter<T, A> {
10071010
IntoIter {
1008-
inner: self.table.into_iter(),
1011+
inner: self.raw.into_iter(),
10091012
}
10101013
}
10111014
}
@@ -1040,7 +1043,7 @@ where
10401043
{
10411044
fn default() -> Self {
10421045
Self {
1043-
table: Default::default(),
1046+
raw: Default::default(),
10441047
}
10451048
}
10461049
}
@@ -1052,7 +1055,7 @@ where
10521055
{
10531056
fn clone(&self) -> Self {
10541057
Self {
1055-
table: self.table.clone(),
1058+
raw: self.raw.clone(),
10561059
}
10571060
}
10581061
}
@@ -1429,7 +1432,7 @@ where
14291432
{
14301433
hash: u64,
14311434
bucket: Bucket<T>,
1432-
table: &'a mut RawTable<T, A>,
1435+
table: &'a mut HashTable<T, A>,
14331436
}
14341437

14351438
unsafe impl<T, A> Send for OccupiedEntry<'_, T, A>
@@ -1496,7 +1499,7 @@ where
14961499
/// # }
14971500
/// ```
14981501
pub fn remove(self) -> (T, VacantEntry<'a, T, A>) {
1499-
let (val, slot) = unsafe { self.table.remove(self.bucket) };
1502+
let (val, slot) = unsafe { self.table.raw.remove(self.bucket) };
15001503
(
15011504
val,
15021505
VacantEntry {
@@ -1642,6 +1645,12 @@ where
16421645
pub fn into_mut(self) -> &'a mut T {
16431646
unsafe { self.bucket.as_mut() }
16441647
}
1648+
1649+
/// Converts the OccupiedEntry into a mutable reference to the underlying
1650+
/// table.
1651+
pub fn into_table(self) -> &'a mut HashTable<T, A> {
1652+
self.table
1653+
}
16451654
}
16461655

16471656
/// A view into a vacant entry in a `HashTable`.
@@ -1689,7 +1698,7 @@ where
16891698
{
16901699
hash: u64,
16911700
insert_slot: InsertSlot,
1692-
table: &'a mut RawTable<T, A>,
1701+
table: &'a mut HashTable<T, A>,
16931702
}
16941703

16951704
impl<T: fmt::Debug, A: Allocator> fmt::Debug for VacantEntry<'_, T, A> {
@@ -1737,6 +1746,7 @@ where
17371746
pub fn insert(self, value: T) -> OccupiedEntry<'a, T, A> {
17381747
let bucket = unsafe {
17391748
self.table
1749+
.raw
17401750
.insert_in_slot(self.hash, self.insert_slot, value)
17411751
};
17421752
OccupiedEntry {
@@ -1745,6 +1755,12 @@ where
17451755
table: self.table,
17461756
}
17471757
}
1758+
1759+
/// Converts the OccupiedEntry into a mutable reference to the underlying
1760+
/// table.
1761+
pub fn into_table(self) -> &'a mut HashTable<T, A> {
1762+
self.table
1763+
}
17481764
}
17491765

17501766
/// An iterator over the entries of a `HashTable` in arbitrary order.

0 commit comments

Comments
 (0)