From 921579cc3cdc57dbcfac48b275d92b7927234988 Mon Sep 17 00:00:00 2001
From: "Chai T. Rex" <ChaiTRex@users.noreply.github.com>
Date: Fri, 10 Apr 2020 12:10:05 -0400
Subject: [PATCH] Add or_insert_with_key to Entry of HashMap/BTreeMap

Going along with or_insert_with, or_insert_with_key provides the
Entry's key to the lambda, avoiding the need to either clone the
key or the need to reimplement this body of this method from
scratch each time.

This is useful when the initial value for a map entry is derived
from the key. For example, the introductory Rust book has an
example Cacher struct that takes an expensive-to-compute lambda and
then can, given an argument to the lambda, produce either the
cached result or execute the lambda.
---
 src/liballoc/collections/btree/map.rs | 27 +++++++++++++++++++++++++++
 src/libstd/collections/hash/map.rs    | 27 +++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index a1e59b2e6af..650582850ab 100644
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -2378,6 +2378,33 @@ impl<'a, K: Ord, V> Entry<'a, K, V> {
         }
     }
 
+    #[unstable(feature = "or_insert_with_key", issue = "70996")]
+    /// Ensures a value is in the entry by inserting, if empty, the result of the default function,
+    /// which takes the key as its argument, and returns a mutable reference to the value in the
+    /// entry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map: BTreeMap<&str, usize> = BTreeMap::new();
+    ///
+    /// map.entry("poneyland").or_insert_with_key(|key| key.chars().count());
+    ///
+    /// assert_eq!(map["poneyland"], 9);
+    /// ```
+    #[inline]
+    pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V {
+        match self {
+            Occupied(entry) => entry.into_mut(),
+            Vacant(entry) => {
+                let value = default(entry.key());
+                entry.insert(value)
+            }
+        }
+    }
+
     /// Returns a reference to this entry's key.
     ///
     /// # Examples
diff --git a/src/libstd/collections/hash/map.rs b/src/libstd/collections/hash/map.rs
index d1cb8e92d56..277b37dd310 100644
--- a/src/libstd/collections/hash/map.rs
+++ b/src/libstd/collections/hash/map.rs
@@ -1943,6 +1943,33 @@ impl<'a, K, V> Entry<'a, K, V> {
         }
     }
 
+    #[unstable(feature = "or_insert_with_key", issue = "70996")]
+    /// Ensures a value is in the entry by inserting, if empty, the result of the default function,
+    /// which takes the key as its argument, and returns a mutable reference to the value in the
+    /// entry.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::collections::HashMap;
+    ///
+    /// let mut map: HashMap<&str, usize> = HashMap::new();
+    ///
+    /// map.entry("poneyland").or_insert_with_key(|key| key.chars().count());
+    ///
+    /// assert_eq!(map["poneyland"], 9);
+    /// ```
+    #[inline]
+    pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V {
+        match self {
+            Occupied(entry) => entry.into_mut(),
+            Vacant(entry) => {
+                let value = default(entry.key());
+                entry.insert(value)
+            }
+        }
+    }
+
     /// Returns a reference to this entry's key.
     ///
     /// # Examples