From 5b9c757fcb81b693f402e3e89deff7a6ef778130 Mon Sep 17 00:00:00 2001
From: Brad Dunbar <dunbarb2@gmail.com>
Date: Tue, 17 Dec 2024 16:47:54 -0500
Subject: [PATCH] HashMap: #or_default_entry and #or_insert_entry

I often find myself in a situation where I would like to insert a key
into a map via `#entry_ref` and then use the key from the resulting
entry.

```
// Get a byte slice from a buffer
let key = client.request.get(index);

// Insert the value (default)
let mut entry = match queues.entry_ref(&key) {
    EntryRef::Occupied(entry) => entry,
    EntryRef::Vacant(entry) => entry.insert_entry(Default::default()),
};

// Use the value
entry.get_mut().insert_back(client.id);

// Reuse the key instead of copying the bytes again
keys.insert(client.id, entry.key());
```

This is common enough that I'd love to have functions for it, similar to
`insert_entry`.

```
// Get a byte slice from a buffer
let key = client.request.get(index);

// Insert the value (default)
let mut entry = queues.entry_ref(&key).or_default_entry();

// Use the value
entry.get_mut().insert_back(client.id);

// Reuse the key instead of copying the bytes again
keys.insert(client.id, entry.key());
```
---
 src/map.rs | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/src/map.rs b/src/map.rs
index c373d5958..e7ea8b9da 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -3527,6 +3527,38 @@ impl<'a, K, V, S, A: Allocator> Entry<'a, K, V, S, A> {
         }
     }
 
+    /// Ensures a value is in the entry by inserting the default if empty,
+    /// and returns an [`OccupiedEntry`].
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use hashbrown::HashMap;
+    ///
+    /// let mut map: HashMap<&str, u32> = HashMap::new();
+    ///
+    /// // nonexistent key
+    /// let entry = map.entry("poneyland").or_insert_entry(3);
+    /// assert_eq!(entry.key(), &"poneyland");
+    /// assert_eq!(entry.get(), &3);
+    ///
+    /// // existing key
+    /// let mut entry = map.entry("poneyland").or_insert_entry(10);
+    /// assert_eq!(entry.key(), &"poneyland");
+    /// assert_eq!(entry.get(), &3);
+    /// ```
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn or_insert_entry(self, default: V) -> OccupiedEntry<'a, K, V, S, A>
+    where
+        K: Hash,
+        S: BuildHasher,
+    {
+        match self {
+            Entry::Occupied(entry) => entry,
+            Entry::Vacant(entry) => entry.insert_entry(default),
+        }
+    }
+
     /// Ensures a value is in the entry by inserting the result of the default function if empty,
     /// and returns a mutable reference to the value in the entry.
     ///
@@ -4328,6 +4360,39 @@ impl<'a, 'b, K, Q: ?Sized, V: Default, S, A: Allocator> EntryRef<'a, 'b, K, Q, V
             EntryRef::Vacant(entry) => entry.insert(Default::default()),
         }
     }
+
+    /// Ensures a value is in the entry by inserting the default value if empty,
+    /// and returns an [`OccupiedEntry`].
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use hashbrown::HashMap;
+    ///
+    /// let mut map: HashMap<String, Option<u32>> = HashMap::new();
+    ///
+    /// // nonexistent key
+    /// let entry = map.entry_ref("poneyland").or_default_entry();
+    /// assert_eq!(entry.key(), &"poneyland");
+    /// assert_eq!(entry.get(), &None);
+    ///
+    /// // existing key
+    /// map.insert("horseland".to_string(), Some(3));
+    /// let entry = map.entry_ref("horseland").or_default_entry();
+    /// assert_eq!(entry.key(), &"horseland");
+    /// assert_eq!(entry.get(), &Some(3));
+    /// ```
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn or_default_entry(self) -> OccupiedEntry<'a, K, V, S, A>
+    where
+        K: Hash + From<&'b Q>,
+        S: BuildHasher,
+    {
+        match self {
+            EntryRef::Occupied(entry) => entry,
+            EntryRef::Vacant(entry) => entry.insert_entry(Default::default()),
+        }
+    }
 }
 
 impl<'a, 'b, K, Q: ?Sized, V, S, A: Allocator> VacantEntryRef<'a, 'b, K, Q, V, S, A> {