Clean up RawTable API

Author: Amanieu

src/map.rs

@@ -1783,19 +1783,16 @@ where
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn insert(&mut self, k: K, v: V) -> Option<V> {
         let hash = make_insert_hash::<K, S>(&self.hash_builder, &k);
-        self.table
-            .reserve(1, make_hasher::<_, V, S>(&self.hash_builder));
-
-        unsafe {
-            let (index, found) = self.table.find_potential(hash, equivalent_key(&k));
-
-            let bucket = self.table.bucket(index);
-
-            if found {
-                Some(mem::replace(&mut bucket.as_mut().1, v))
-            } else {
-                self.table.mark_inserted(index, hash);
-                bucket.write((k, v));
+        let hasher = make_hasher::<_, V, S>(&self.hash_builder);
+        match self
+            .table
+            .find_or_find_insert_slot(hash, equivalent_key(&k), hasher)
+        {
+            Ok(bucket) => Some(mem::replace(unsafe { &mut bucket.as_mut().1 }, v)),
+            Err(slot) => {
+                unsafe {
+                    self.table.insert_in_slot(hash, slot, (k, v));
+                }
                 None
             }
         }
@@ -2178,7 +2175,7 @@ impl<K, V, S, A: Allocator + Clone> HashMap<K, V, S, A> {
     /// {
     ///     let raw_table = map.raw_table_mut();
     ///     match raw_table.find(hash, is_match) {
-    ///         Some(bucket) => Some(unsafe { raw_table.remove(bucket) }),
+    ///         Some(bucket) => Some(unsafe { raw_table.remove(bucket).0 }),
     ///         None => None,
     ///     }
     /// }
@@ -2801,7 +2798,7 @@ impl<K, V, A: Allocator + Clone> ExtractIfInner<'_, K, V, A> {
             for item in &mut self.iter {
                 let &mut (ref key, ref mut value) = item.as_mut();
                 if f(key, value) {
-                    return Some(self.table.remove(item));
+                    return Some(self.table.remove(item).0);
                 }
             }
         }
@@ -3922,7 +3919,7 @@ impl<'a, K, V, S, A: Allocator + Clone> RawOccupiedEntryMut<'a, K, V, S, A> {
     /// ```
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn remove_entry(self) -> (K, V) {
-        unsafe { self.table.remove(self.elem) }
+        unsafe { self.table.remove(self.elem).0 }
     }
 
     /// Provides shared access to the key and owned access to the value of
@@ -5295,7 +5292,7 @@ impl<'a, K, V, S, A: Allocator + Clone> OccupiedEntry<'a, K, V, S, A> {
     /// ```
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn remove_entry(self) -> (K, V) {
-        unsafe { self.table.table.remove(self.elem) }
+        unsafe { self.table.table.remove(self.elem).0 }
     }
 
     /// Gets a reference to the value in the entry.
@@ -6017,7 +6014,7 @@ impl<'a, 'b, K, Q: ?Sized, V, S, A: Allocator + Clone> OccupiedEntryRef<'a, 'b,
     /// ```
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn remove_entry(self) -> (K, V) {
-        unsafe { self.table.table.remove(self.elem) }
+        unsafe { self.table.table.remove(self.elem).0 }
     }
 
     /// Gets a reference to the value in the entry.
@@ -8365,7 +8362,7 @@ mod test_map {
                         let e = map.table.find(hash_value, |q| q.0.eq(&i));
                         if let Some(e) = e {
                             it.reflect_remove(&e);
-                            let t = map.table.remove(e);
+                            let t = map.table.remove(e).0;
                             removed.push(t);
                             left -= 1;
                         } else {

src/raw/mod.rs

@@ -282,6 +282,11 @@ impl TableLayout {
     }
 }
 
+/// A reference to an empty bucket into which an can be inserted.
+pub struct InsertSlot {
+    index: usize,
+}
+
 /// A reference to a hash table bucket containing a `T`.
 ///
 /// This is usually just a pointer to the element itself. However if the element
@@ -1001,19 +1006,26 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
     }
 
     /// Removes an element from the table, returning it.
+    ///
+    /// This also returns an `InsertSlot` pointing to the newly free bucket.
     #[cfg_attr(feature = "inline-more", inline)]
     #[allow(clippy::needless_pass_by_value)]
-    pub unsafe fn remove(&mut self, item: Bucket<T>) -> T {
+    pub unsafe fn remove(&mut self, item: Bucket<T>) -> (T, InsertSlot) {
         self.erase_no_drop(&item);
-        item.read()
+        (
+            item.read(),
+            InsertSlot {
+                index: self.bucket_index(&item),
+            },
+        )
     }
 
     /// Finds and removes an element from the table, returning it.
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn remove_entry(&mut self, hash: u64, eq: impl FnMut(&T) -> bool) -> Option<T> {
         // Avoid `Option::map` because it bloats LLVM IR.
         match self.find(hash, eq) {
-            Some(bucket) => Some(unsafe { self.remove(bucket) }),
+            Some(bucket) => Some(unsafe { self.remove(bucket).0 }),
             None => None,
         }
     }
@@ -1161,22 +1173,18 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn insert(&mut self, hash: u64, value: T, hasher: impl Fn(&T) -> u64) -> Bucket<T> {
         unsafe {
-            let mut index = self.table.find_insert_slot(hash);
+            let mut slot = self.table.find_insert_slot(hash);
 
             // We can avoid growing the table once we have reached our load
             // factor if we are replacing a tombstone. This works since the
             // number of EMPTY slots does not change in this case.
-            let old_ctrl = *self.table.ctrl(index);
+            let old_ctrl = *self.table.ctrl(slot.index);
             if unlikely(self.table.growth_left == 0 && special_is_empty(old_ctrl)) {
                 self.reserve(1, hasher);
-                index = self.table.find_insert_slot(hash);
+                slot = self.table.find_insert_slot(hash);
             }
 
-            self.table.record_item_insert_at(index, old_ctrl, hash);
-
-            let bucket = self.bucket(index);
-            bucket.write(value);
-            bucket
+            self.insert_in_slot(hash, slot, value)
         }
     }
 
@@ -1244,7 +1252,7 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
         let old_ctrl = *self.table.ctrl(index);
         debug_assert!(self.is_bucket_full(index));
         let old_growth_left = self.table.growth_left;
-        let item = self.remove(bucket);
+        let item = self.remove(bucket).0;
         if let Some(new_item) = f(item) {
             self.table.growth_left = old_growth_left;
             self.table.set_ctrl(index, old_ctrl);
@@ -1256,20 +1264,47 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
         }
     }
 
-    /// Searches for an element in the table,
-    /// or a potential slot where that element could be inserted.
+    /// Searches for an element in the table. If the element is not found,
+    /// returns `Err` with the position of a slot where an element with the
+    /// same hash could be inserted.
+    ///
+    /// This function may resize the table if additional space is required for
+    /// inserting an element.
     #[inline]
-    pub fn find_potential(&self, hash: u64, mut eq: impl FnMut(&T) -> bool) -> (usize, bool) {
-        self.table.find_potential_inner(hash, &mut |index| unsafe {
-            eq(self.bucket(index).as_ref())
-        })
+    pub fn find_or_find_insert_slot(
+        &mut self,
+        hash: u64,
+        mut eq: impl FnMut(&T) -> bool,
+        hasher: impl Fn(&T) -> u64,
+    ) -> Result<Bucket<T>, InsertSlot> {
+        self.reserve(1, hasher);
+
+        match self
+            .table
+            .find_or_find_insert_slot_inner(hash, &mut |index| unsafe {
+                eq(self.bucket(index).as_ref())
+            }) {
+            Ok(index) => Ok(unsafe { self.bucket(index) }),
+            Err(slot) => Err(slot),
+        }
     }
 
-    /// Marks an element in the table as inserted.
+    /// Inserts a new element into the table in the given slot, and returns its
+    /// raw bucket.
+    ///
+    /// # Safety
+    ///
+    /// `slot` must point to a slot previously returned by
+    /// `find_or_find_insert_slot`, and no mutation of the table must have
+    /// occurred since that call.
     #[inline]
-    pub unsafe fn mark_inserted(&mut self, index: usize, hash: u64) {
-        let old_ctrl = *self.table.ctrl(index);
-        self.table.record_item_insert_at(index, old_ctrl, hash);
+    pub unsafe fn insert_in_slot(&mut self, hash: u64, slot: InsertSlot, value: T) -> Bucket<T> {
+        let old_ctrl = *self.table.ctrl(slot.index);
+        self.table.record_item_insert_at(slot.index, old_ctrl, hash);
+
+        let bucket = self.bucket(slot.index);
+        bucket.write(value);
+        bucket
     }
 
     /// Searches for an element in the table.
@@ -1608,7 +1643,7 @@ impl<A: Allocator + Clone> RawTableInner<A> {
     /// Fixes up an insertion slot due to false positives for groups smaller than the group width.
     /// This must only be used on insertion slots found by `find_insert_slot_in_group`.
     #[inline]
-    unsafe fn fix_insert_slot(&self, index: usize) -> usize {
+    unsafe fn fix_insert_slot(&self, mut index: usize) -> InsertSlot {
         // In tables smaller than the group width
         // (self.buckets() < Group::WIDTH), trailing control
         // bytes outside the range of the table are filled with
@@ -1636,12 +1671,11 @@ impl<A: Allocator + Clone> RawTableInner<A> {
             //   with EMPTY bytes, so this second scan either finds an empty slot (due to the
             //   load factor) or hits the trailing control bytes (containing EMPTY). See
             //   `intrinsics::cttz_nonzero` for more information.
-            Group::load_aligned(self.ctrl(0))
+            index = Group::load_aligned(self.ctrl(0))
                 .match_empty_or_deleted()
-                .lowest_set_bit_nonzero()
-        } else {
-            index
+                .lowest_set_bit_nonzero();
         }
+        InsertSlot { index }
     }
 
     /// Finds the position to insert something in a group.
@@ -1663,11 +1697,11 @@ impl<A: Allocator + Clone> RawTableInner<A> {
     /// This uses dynamic dispatch to reduce the amount of code generated, but that is
     /// eliminated by LLVM optimizations.
     #[inline]
-    pub fn find_potential_inner(
+    fn find_or_find_insert_slot_inner(
         &self,
         hash: u64,
         eq: &mut dyn FnMut(usize) -> bool,
-    ) -> (usize, bool) {
+    ) -> Result<usize, InsertSlot> {
         let mut insert_slot = None;
 
         let h2_hash = h2(hash);
@@ -1680,7 +1714,7 @@ impl<A: Allocator + Clone> RawTableInner<A> {
                 let index = (probe_seq.pos + bit) & self.bucket_mask;
 
                 if likely(eq(index)) {
-                    return (index, true);
+                    return Ok(index);
                 }
             }
 
@@ -1697,7 +1731,7 @@ impl<A: Allocator + Clone> RawTableInner<A> {
                 // least one. For tables smaller than the group width, there will still be an
                 // empty element in the current (and only) group due to the load factor.
                 unsafe {
-                    return (self.fix_insert_slot(insert_slot.unwrap_unchecked()), false);
+                    return Err(self.fix_insert_slot(insert_slot.unwrap_unchecked()));
                 }
             }
 
@@ -1711,7 +1745,7 @@ impl<A: Allocator + Clone> RawTableInner<A> {
     /// There must be at least 1 empty bucket in the table.
     #[inline]
     unsafe fn prepare_insert_slot(&self, hash: u64) -> (usize, u8) {
-        let index = self.find_insert_slot(hash);
+        let index = self.find_insert_slot(hash).index;
         let old_ctrl = *self.ctrl(index);
         self.set_ctrl_h2(index, hash);
         (index, old_ctrl)
@@ -1739,7 +1773,7 @@ impl<A: Allocator + Clone> RawTableInner<A> {
     ///
     /// [`undefined behavior`]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
     #[inline]
-    fn find_insert_slot(&self, hash: u64) -> usize {
+    fn find_insert_slot(&self, hash: u64) -> InsertSlot {
         let mut probe_seq = self.probe_seq(hash);
         loop {
             // SAFETY:
@@ -1922,7 +1956,7 @@ impl<A: Allocator + Clone> RawTableInner<A> {
     #[cfg(feature = "raw")]
     #[inline]
     unsafe fn prepare_insert_no_grow(&mut self, hash: u64) -> Result<usize, ()> {
-        let index = self.find_insert_slot(hash);
+        let index = self.find_insert_slot(hash).index;
         let old_ctrl = *self.ctrl(index);
         if unlikely(self.growth_left == 0 && special_is_empty(old_ctrl)) {
             Err(())
@@ -2293,7 +2327,7 @@ impl<A: Allocator + Clone> RawTableInner<A> {
                 let hash = hasher(*guard, i);
 
                 // Search for a suitable place to put it
-                let new_i = guard.find_insert_slot(hash);
+                let new_i = guard.find_insert_slot(hash).index;
 
                 // Probing works by scanning through all of the control
                 // bytes in groups, which may not be aligned to the group

src/rustc_entry.rs

@@ -330,7 +330,7 @@ impl<'a, K, V, A: Allocator + Clone> RustcOccupiedEntry<'a, K, V, A> {
     /// ```
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn remove_entry(self) -> (K, V) {
-        unsafe { self.table.remove(self.elem) }
+        unsafe { self.table.remove(self.elem).0 }
     }
 
     /// Gets a reference to the value in the entry.