Add hash_table::OccupiedEntry::replace_entry_with to mirror HashMap API

Author: clarfonthey

src/raw/mod.rs

@@ -1127,11 +1127,16 @@ impl<T, A: Allocator> RawTable<T, A> {
     /// Temporary removes a bucket, applying the given function to the removed
     /// element and optionally put back the returned value in the same bucket.
     ///
-    /// Returns `true` if the bucket still contains an element
+    /// Returns `true` if the bucket still contains an element. Always returns the
+    /// associated bucket index and tag.
     ///
     /// This does not check if the given bucket is actually occupied.
     #[cfg_attr(feature = "inline-more", inline)]
-    pub unsafe fn replace_bucket_with<F>(&mut self, bucket: Bucket<T>, f: F) -> bool
+    pub(crate) unsafe fn replace_bucket_with<F>(
+        &mut self,
+        bucket: Bucket<T>,
+        f: F,
+    ) -> (bool, usize, Tag)
     where
         F: FnOnce(T) -> Option<T>,
     {
@@ -1140,15 +1145,19 @@ impl<T, A: Allocator> RawTable<T, A> {
         debug_assert!(self.is_bucket_full(index));
         let old_growth_left = self.table.growth_left;
         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);
-            self.table.items += 1;
-            self.bucket(index).write(new_item);
-            true
-        } else {
-            false
-        }
+        (
+            if let Some(new_item) = f(item) {
+                self.table.growth_left = old_growth_left;
+                self.table.set_ctrl(index, old_ctrl);
+                self.table.items += 1;
+                self.bucket(index).write(new_item);
+                true
+            } else {
+                false
+            },
+            index,
+            old_ctrl,
+        )
     }
 
     /// Searches for an element in the table. If the element is not found,

src/raw_entry.rs

@@ -1282,7 +1282,7 @@ impl<'a, K, V, S, A: Allocator> RawOccupiedEntryMut<'a, K, V, S, A> {
         F: FnOnce(&K, V) -> Option<V>,
     {
         unsafe {
-            let still_occupied = self
+            let (still_occupied, _, _) = self
                 .table
                 .replace_bucket_with(self.elem.clone(), |(key, value)| {
                     f(&key, value).map(|new_value| (key, new_value))

src/table.rs

@@ -2266,6 +2266,79 @@ where
     pub fn bucket_index(&self) -> usize {
         unsafe { self.table.raw.bucket_index(&self.bucket) }
     }
+
+    /// Provides owned access to the value of the entry and allows to replace or
+    /// remove it based on the value of the returned option.
+    ///
+    /// The hash of the new item should be the same as the old item.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[cfg(feature = "nightly")]
+    /// # fn test() {
+    /// use hashbrown::{HashTable, DefaultHashBuilder};
+    /// use hashbrown::hash_table::Entry;
+    /// use std::hash::BuildHasher;
+    ///
+    /// let mut table = HashTable::new();
+    /// let hasher = DefaultHashBuilder::default();
+    /// let hasher = |(ref key, _): &_| hasher.hash_one(key);
+    /// table.insert_unique(hasher(&("poneyland", 42)), ("poneyland", 42), hasher);
+    ///
+    /// let entry = match table.entry(hasher(&("poneyland", 42)), |entry| entry.0 == "poneyland", hasher) {
+    ///     Entry::Occupied(e) => unsafe {
+    ///         e.replace_entry_with(|(k, v)| {
+    ///             assert_eq!(k, "poneyland");
+    ///             assert_eq!(v, 42);
+    ///             Some(("poneyland", v + 1))
+    ///         })
+    ///     }
+    ///     Entry::Vacant(_) => panic!(),
+    /// };
+    ///
+    /// match entry {
+    ///     Entry::Occupied(e) => {
+    ///         assert_eq!(e.get(), &("poneyland", 43));
+    ///     }
+    ///     Entry::Vacant(_) => panic!(),
+    /// }
+    ///
+    /// let entry = match table.entry(hasher(&("poneyland", 43)), |entry| entry.0 == "poneyland", hasher) {
+    ///     Entry::Occupied(e) => unsafe { e.replace_entry_with(|(_k, _v)| None) },
+    ///     Entry::Vacant(_) => panic!(),
+    /// };
+    ///
+    /// match entry {
+    ///     Entry::Vacant(e) => {
+    ///         // nice!
+    ///     }
+    ///     Entry::Occupied(_) => panic!(),
+    /// }
+    ///
+    /// assert!(table.is_empty());
+    /// # }
+    /// # fn main() {
+    /// #     #[cfg(feature = "nightly")]
+    /// #     test()
+    /// # }
+    /// ```
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn replace_entry_with<F>(self, f: F) -> Entry<'a, T, A>
+    where
+        F: FnOnce(T) -> Option<T>,
+    {
+        unsafe {
+            match self.table.raw.replace_bucket_with(self.bucket.clone(), f) {
+                (true, _, _) => Entry::Occupied(self),
+                (false, index, tag) => Entry::Vacant(VacantEntry {
+                    tag,
+                    index,
+                    table: self.table,
+                }),
+            }
+        }
+    }
 }
 
 /// A view into a vacant entry in a `HashTable`.