Extend map with referring insertion by nikita-volkov · Pull Request #60142 · rust-lang/rust (original) (raw)

Implement (HashMap) Entry::insert as per rust-lang#60142

Implementation of Entry::insert as per @SimonSapin's comment on rust-lang#60142. This requires a patch to hashbrown:

diff --git a/src/rustc_entry.rs b/src/rustc_entry.rs
index fefa5c3..7de8300 100644
--- a/src/rustc_entry.rs
+++ b/src/rustc_entry.rs
@@ -546,6 +546,32 @@ impl<'a, K, V> RustcVacantEntry<'a, K, V> {
         let bucket = self.table.insert_no_grow(self.hash, (self.key, value));
         unsafe { &mut bucket.as_mut().1 }
     }
+
+    /// Sets the value of the entry with the RustcVacantEntry's key,
+    /// and returns a RustcOccupiedEntry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use hashbrown::HashMap;
+    /// use hashbrown::hash_map::RustcEntry;
+    ///
+    /// let mut map: HashMap<&str, u32> = HashMap::new();
+    ///
+    /// if let RustcEntry::Vacant(v) = map.rustc_entry("poneyland") {
+    ///     let o = v.insert_and_return(37);
+    ///     assert_eq!(o.get(), &37);
+    /// }
+    /// ```
+     #[inline]
+    pub fn insert_and_return(self, value: V) -> RustcOccupiedEntry<'a, K, V> {
+        let bucket = self.table.insert_no_grow(self.hash, (self.key, value));
+        RustcOccupiedEntry {
+            key: None,
+            elem: bucket,
+            table: self.table
+        }
+    }
 }

 impl<K, V> IterMut<'_, K, V> {

This is also only an implementation for HashMap. I tried implementing for BTreeMap, but I don't really understand BTreeMap's internals and require more guidance on implementing the equivalent VacantEntry::insert_and_return such that it returns an OccupiedEntry. Notably, following the original PR's modifications I end up needing a Handle<NodeRef<marker::Mut<'_>, _, _, marker::LeafOrInternal>, _> while I only have a Handle<NodeRef<marker::Mut<'_>, _, _, marker::Internal>, _> and don't know how to proceed.

(To be clear, I'm not asking for guidance right now; I'd be happy getting only the HashMap implementation — the subject of this PR — reviewed and ready, and leave the BTreeMap implementation for a latter PR.)