Rollup merge of #133042 - cuviper:btreemap-insert_entry, r=Amanieu

btree: add `{Entry,VacantEntry}::insert_entry`

This matches the recently-stabilized methods on `HashMap` entries. I've
reused tracking issue #65225 for now, but we may want to split it.

2 files changed, 76 insertions, 31 deletions

diff --git a/library/alloc/src/collections/btree/map/entry.rs b/library/alloc/src/collections/btree/map/entry.rs
index 75bb86916a8..0da6af54bc2 100644
--- a/library/alloc/src/collections/btree/map/entry.rs
+++ b/library/alloc/src/collections/btree/map/entry.rs
@@ -269,6 +269,31 @@ impl<'a, K: Ord, V, A: Allocator + Clone> Entry<'a, K, V, A> {
             Vacant(entry) => Vacant(entry),
         }
     }
+
+    /// Sets the value of the entry, and returns an `OccupiedEntry`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_entry_insert)]
+    /// use std::collections::BTreeMap;
+    ///
+    /// let mut map: BTreeMap<&str, String> = BTreeMap::new();
+    /// let entry = map.entry("poneyland").insert_entry("hoho".to_string());
+    ///
+    /// assert_eq!(entry.key(), &"poneyland");
+    /// ```
+    #[inline]
+    #[unstable(feature = "btree_entry_insert", issue = "65225")]
+    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V, A> {
+        match self {
+            Occupied(mut entry) => {
+                entry.insert(value);
+                entry
+            }
+            Vacant(entry) => entry.insert_entry(value),
+        }
+    }
 }
 
 impl<'a, K: Ord, V: Default, A: Allocator + Clone> Entry<'a, K, V, A> {
@@ -348,41 +373,61 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     #[rustc_confusables("push", "put")]
-    pub fn insert(mut self, value: V) -> &'a mut V {
-        let out_ptr = match self.handle {
+    pub fn insert(self, value: V) -> &'a mut V {
+        self.insert_entry(value).into_mut()
+    }
+
+    /// Sets the value of the entry with the `VacantEntry`'s key,
+    /// and returns an `OccupiedEntry`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(btree_entry_insert)]
+    /// use std::collections::BTreeMap;
+    /// use std::collections::btree_map::Entry;
+    ///
+    /// let mut map: BTreeMap<&str, u32> = BTreeMap::new();
+    ///
+    /// if let Entry::Vacant(o) = map.entry("poneyland") {
+    ///     let entry = o.insert_entry(37);
+    ///     assert_eq!(entry.get(), &37);
+    /// }
+    /// assert_eq!(map["poneyland"], 37);
+    /// ```
+    #[unstable(feature = "btree_entry_insert", issue = "65225")]
+    pub fn insert_entry(mut self, value: V) -> OccupiedEntry<'a, K, V, A> {
+        let handle = match self.handle {
             None => {
                 // SAFETY: There is no tree yet so no reference to it exists.
-                let map = unsafe { self.dormant_map.awaken() };
-                let mut root = NodeRef::new_leaf(self.alloc.clone());
-                let val_ptr = root.borrow_mut().push(self.key, value);
-                map.root = Some(root.forget_type());
-                map.length = 1;
-                val_ptr
-            }
-            Some(handle) => {
-                let new_handle =
-                    handle.insert_recursing(self.key, value, self.alloc.clone(), |ins| {
-                        drop(ins.left);
-                        // SAFETY: Pushing a new root node doesn't invalidate
-                        // handles to existing nodes.
-                        let map = unsafe { self.dormant_map.reborrow() };
-                        let root = map.root.as_mut().unwrap(); // same as ins.left
-                        root.push_internal_level(self.alloc).push(ins.kv.0, ins.kv.1, ins.right)
-                    });
-
-                // Get the pointer to the value
-                let val_ptr = new_handle.into_val_mut();
-
-                // SAFETY: We have consumed self.handle.
-                let map = unsafe { self.dormant_map.awaken() };
-                map.length += 1;
-                val_ptr
+                let map = unsafe { self.dormant_map.reborrow() };
+                let root = map.root.insert(NodeRef::new_leaf(self.alloc.clone()).forget_type());
+                // SAFETY: We *just* created the root as a leaf, and we're
+                // stacking the new handle on the original borrow lifetime.
+                unsafe {
+                    let mut leaf = root.borrow_mut().cast_to_leaf_unchecked();
+                    leaf.push_with_handle(self.key, value)
+                }
             }
+            Some(handle) => handle.insert_recursing(self.key, value, self.alloc.clone(), |ins| {
+                drop(ins.left);
+                // SAFETY: Pushing a new root node doesn't invalidate
+                // handles to existing nodes.
+                let map = unsafe { self.dormant_map.reborrow() };
+                let root = map.root.as_mut().unwrap(); // same as ins.left
+                root.push_internal_level(self.alloc.clone()).push(ins.kv.0, ins.kv.1, ins.right)
+            }),
         };
 
-        // Now that we have finished growing the tree using borrowed references,
-        // dereference the pointer to a part of it, that we picked up along the way.
-        unsafe { &mut *out_ptr }
+        // SAFETY: modifying the length doesn't invalidate handles to existing nodes.
+        unsafe { self.dormant_map.reborrow().length += 1 };
+
+        OccupiedEntry {
+            handle: handle.forget_node_type(),
+            dormant_map: self.dormant_map,
+            alloc: self.alloc,
+            _marker: PhantomData,
+        }
     }
 }
 
diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index 64a13bb6a0b..0c93eff0d20 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -739,7 +739,7 @@ impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
 
 impl<'a, K, V> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
     /// Unsafely asserts to the compiler the static information that this node is a `Leaf`.
-    unsafe fn cast_to_leaf_unchecked(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
+    pub unsafe fn cast_to_leaf_unchecked(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
         debug_assert!(self.height == 0);
         NodeRef { height: self.height, node: self.node, _marker: PhantomData }
     }