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

Author: clarfonthey

src/raw/mod.rs

@@ -1124,14 +1124,14 @@ impl<T, A: Allocator> RawTable<T, A> {
         bucket
     }
 
-    /// Temporary removes a bucket, applying the given function to the removed
+    /// Temporarily 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 tag for bucket if the bucket is emptied out.
     ///
     /// 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) -> Option<Tag>
     where
         F: FnOnce(T) -> Option<T>,
     {
@@ -1145,9 +1145,9 @@ impl<T, A: Allocator> RawTable<T, A> {
             self.table.set_ctrl(index, old_ctrl);
             self.table.items += 1;
             self.bucket(index).write(new_item);
-            true
+            None
         } else {
-            false
+            Some(old_ctrl)
         }
     }
 

src/raw_entry.rs

@@ -1282,13 +1282,13 @@ 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 tag = self
                 .table
                 .replace_bucket_with(self.elem.clone(), |(key, value)| {
                     f(&key, value).map(|new_value| (key, new_value))
                 });
 
-            if still_occupied {
+            if tag.is_none() {
                 RawEntryMut::Occupied(self)
             } else {
                 RawEntryMut::Vacant(RawVacantEntryMut {

src/table.rs

@@ -2246,6 +2246,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) {
+                None => Entry::Occupied(self),
+                Some(tag) => Entry::Vacant(VacantEntry {
+                    tag,
+                    index: self.bucket_index(),
+                    table: self.table,
+                }),
+            }
+        }
+    }
 }
 
 /// A view into a vacant entry in a `HashTable`.