feat: add HashTable::insert_unique_within_capacity

Author: morrisonlevi

src/raw/mod.rs

@@ -1045,18 +1045,15 @@ impl<T, A: Allocator> RawTable<T, A> {
         )
     }
 
-    /// Inserts a new element into the table, and returns its raw bucket.
-    ///
-    /// This does not check if the given element already exists in the table.
-    #[cfg_attr(feature = "inline-more", inline)]
-    pub fn insert(&mut self, hash: u64, value: T, hasher: impl Fn(&T) -> u64) -> Bucket<T> {
+    #[inline(always)]
+    fn find_insert_slot(&self, hash: u64) -> Option<InsertSlot> {
         unsafe {
             // SAFETY:
             // 1. The [`RawTableInner`] must already have properly initialized control bytes since
             //    we will never expose `RawTable::new_uninitialized` in a public API.
             //
             // 2. We reserve additional space (if necessary) right after calling this function.
-            let mut slot = self.table.find_insert_slot(hash);
+            let 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.
@@ -1065,14 +1062,40 @@ impl<T, A: Allocator> RawTable<T, A> {
             // in the range `0..=self.buckets()`.
             let old_ctrl = *self.table.ctrl(slot.index);
             if unlikely(self.table.growth_left == 0 && old_ctrl.special_is_empty()) {
+                None
+            } else {
+                Some(slot)
+            }
+        }
+    }
+
+    /// Inserts a new element into the table, and returns its raw bucket.
+    ///
+    /// This does not check if the given element already exists in the table.
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn insert(&mut self, hash: u64, value: T, hasher: impl Fn(&T) -> u64) -> Bucket<T> {
+        let slot = match self.find_insert_slot(hash) {
+            None => {
                 self.reserve(1, hasher);
                 // SAFETY: We know for sure that `RawTableInner` has control bytes
                 // initialized and that there is extra space in the table.
-                slot = self.table.find_insert_slot(hash);
+                unsafe { self.table.find_insert_slot(hash) }
             }
+            Some(slot) => slot,
+        };
+        // SAFETY: todo
+        unsafe { self.insert_in_slot(hash, slot, value) }
+    }
 
-            self.insert_in_slot(hash, slot, value)
-        }
+    /// Inserts a new element into the table if there is capacity, and returns
+    /// its raw bucket.
+    ///
+    /// This does not check if the given element already exists in the table.
+    #[inline]
+    pub(crate) fn insert_within_capacity(&mut self, hash: u64, value: T) -> Option<Bucket<T>> {
+        let slot = self.find_insert_slot(hash)?;
+        // SAFETY: todo
+        Some(unsafe { self.insert_in_slot(hash, slot, value) })
     }
 
     /// Inserts a new element into the table, and returns a mutable reference to it.

src/table.rs

@@ -415,6 +415,44 @@ where
         }
     }
 
+    /// Inserts an element into the `HashTable` with the given hash value, but
+    /// without checking whether an equivalent element already exists within
+    /// the table. If there is insufficient capacity, then this returns None.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[cfg(feature = "nightly")]
+    /// # fn test() {
+    /// use hashbrown::{HashTable, DefaultHashBuilder};
+    /// use std::hash::BuildHasher;
+    ///
+    /// let mut v = HashTable::new();
+    /// let hasher = DefaultHashBuilder::default();
+    /// let hasher = |val: &_| hasher.hash_one(val);
+    /// assert_eq!(None, v.insert_unique_within_capacity(hasher(&1), 1));
+    /// v.reserve(1, hasher);
+    /// v.insert_unique_within_capacity(hasher(&1), 1);
+    /// assert!(!v.is_empty());
+    /// # }
+    /// # fn main() {
+    /// #     #[cfg(feature = "nightly")]
+    /// #     test()
+    /// # }
+    /// ```
+    pub fn insert_unique_within_capacity(
+        &mut self,
+        hash: u64,
+        value: T,
+    ) -> Option<OccupiedEntry<'_, T, A>> {
+        let bucket = self.raw.insert_within_capacity(hash, value)?;
+        Some(OccupiedEntry {
+            hash,
+            bucket,
+            table: self,
+        })
+    }
+
     /// Clears the table, removing all values.
     ///
     /// # Examples