BTreeMap::entry: Avoid allocating if no insertion

library/alloc/src/collections/btree/map.rs

@@ -212,7 +212,7 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
                     let mut out_tree = clone_subtree(internal.first_edge().descend());
 
                     {
-                        let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
+                        let out_root = out_tree.root.as_mut().unwrap();
                         let mut out_node = out_root.push_internal_level();
                         let mut in_edge = internal.first_edge();
                         while let Ok(kv) = in_edge.right_kv() {
@@ -278,11 +278,12 @@ where
 
     fn replace(&mut self, key: K) -> Option<K> {
         let (map, dormant_map) = DormantMutRef::new(self);
-        let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
+        let root_node = map.root.get_or_insert_with(Root::new).borrow_mut();
         match root_node.search_tree::<K>(&key) {
             Found(mut kv) => Some(mem::replace(kv.key_mut(), key)),
             GoDown(handle) => {
-                VacantEntry { key, handle, dormant_map, _marker: PhantomData }.insert(());
+                VacantEntry { key, handle: Some(handle), dormant_map, _marker: PhantomData }
+                    .insert(());
                 None
             }
         }
@@ -1032,7 +1033,7 @@ impl<K, V> BTreeMap<K, V> {
 
         let self_iter = mem::take(self).into_iter();
         let other_iter = mem::take(other).into_iter();
-        let root = BTreeMap::ensure_is_owned(&mut self.root);
+        let root = self.root.get_or_insert_with(Root::new);
         root.append_from_sorted_iters(self_iter, other_iter, &mut self.length)
     }
 
@@ -1144,14 +1145,20 @@ impl<K, V> BTreeMap<K, V> {
     where
         K: Ord,
     {
-        // FIXME(@porglezomp) Avoid allocating if we don't insert
         let (map, dormant_map) = DormantMutRef::new(self);
-        let root_node = Self::ensure_is_owned(&mut map.root).borrow_mut();
-        match root_node.search_tree(&key) {
-            Found(handle) => Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData }),
-            GoDown(handle) => {
-                Vacant(VacantEntry { key, handle, dormant_map, _marker: PhantomData })
-            }
+        match map.root {
+            None => Vacant(VacantEntry { key, handle: None, dormant_map, _marker: PhantomData }),
+            Some(ref mut root) => match root.borrow_mut().search_tree(&key) {
+                Found(handle) => {
+                    Occupied(OccupiedEntry { handle, dormant_map, _marker: PhantomData })
+                }
+                GoDown(handle) => Vacant(VacantEntry {
+                    key,
+                    handle: Some(handle),
+                    dormant_map,
+                    _marker: PhantomData,
+                }),
+            },
         }
     }
 
@@ -2247,12 +2254,6 @@ impl<K, V> BTreeMap<K, V> {
     pub const fn is_empty(&self) -> bool {
         self.len() == 0
     }
-
-    /// If the root node is the empty (non-allocated) root node, allocate our
-    /// own node. Is an associated function to avoid borrowing the entire BTreeMap.
-    fn ensure_is_owned(root: &mut Option<Root<K, V>>) -> &mut Root<K, V> {
-        root.get_or_insert_with(Root::new)
-    }
 }
 
 #[cfg(test)]

library/alloc/src/collections/btree/map/entry.rs

@@ -40,7 +40,8 @@ impl<K: Debug + Ord, V: Debug> Debug for Entry<'_, K, V> {
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct VacantEntry<'a, K: 'a, V: 'a> {
     pub(super) key: K,
-    pub(super) handle: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
+    /// `None` for a (empty) map without root
+    pub(super) handle: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
     pub(super) dormant_map: DormantMutRef<'a, BTreeMap<K, V>>,
 
     // Be invariant in `K` and `V`
@@ -312,22 +313,33 @@ impl<'a, K: Ord, V> VacantEntry<'a, K, V> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn insert(self, value: V) -> &'a mut V {
-        let out_ptr = match self.handle.insert_recursing(self.key, value) {
-            (None, val_ptr) => {
-                // SAFETY: We have consumed self.handle and the handle returned.
-                let map = unsafe { self.dormant_map.awaken() };
-                map.length += 1;
-                val_ptr
-            }
-            (Some(ins), val_ptr) => {
-                drop(ins.left);
+        let out_ptr = match self.handle {
+            None => {
                 // SAFETY: We have consumed self.handle and the reference returned.
                 let map = unsafe { self.dormant_map.awaken() };
-                let root = map.root.as_mut().unwrap();
-                root.push_internal_level().push(ins.kv.0, ins.kv.1, ins.right);
-                map.length += 1;
+                let mut root = NodeRef::new_leaf();
+                let val_ptr = root.borrow_mut().push(self.key, value) as *mut V;
+                map.root = Some(root.forget_type());
+                map.length = 1;
                 val_ptr
             }
+            Some(handle) => match handle.insert_recursing(self.key, value) {
+                (None, val_ptr) => {
+                    // SAFETY: We have consumed self.handle and the handle returned.
+                    let map = unsafe { self.dormant_map.awaken() };
+                    map.length += 1;
+                    val_ptr
+                }
+                (Some(ins), val_ptr) => {
+                    drop(ins.left);
+                    // SAFETY: We have consumed self.handle and the reference returned.
+                    let map = unsafe { self.dormant_map.awaken() };
+                    let root = map.root.as_mut().unwrap();
+                    root.push_internal_level().push(ins.kv.0, ins.kv.1, ins.right);
+                    map.length += 1;
+                    val_ptr
+                }
+            },
         };
         // 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.

library/alloc/src/collections/btree/map/tests.rs

@@ -115,8 +115,9 @@ impl<K, V> BTreeMap<K, V> {
         K: Ord,
     {
         let iter = mem::take(self).into_iter();
-        let root = BTreeMap::ensure_is_owned(&mut self.root);
-        root.bulk_push(iter, &mut self.length);
+        if !iter.is_empty() {
+            self.root.insert(Root::new()).bulk_push(iter, &mut self.length);
+        }
     }
 }
 
@@ -914,7 +915,7 @@ mod test_drain_filter {
     fn empty() {
         let mut map: BTreeMap<i32, i32> = BTreeMap::new();
         map.drain_filter(|_, _| unreachable!("there's nothing to decide on"));
-        assert!(map.is_empty());
+        assert_eq!(map.height(), None);
         map.check();
     }
 
@@ -1410,7 +1411,7 @@ fn test_clear() {
         assert_eq!(map.len(), len);
         map.clear();
         map.check();
-        assert!(map.is_empty());
+        assert_eq!(map.height(), None);
     }
 }
 
@@ -1789,7 +1790,7 @@ fn test_occupied_entry_key() {
     let mut a = BTreeMap::new();
     let key = "hello there";
     let value = "value goes here";
-    assert!(a.is_empty());
+    assert_eq!(a.height(), None);
     a.insert(key, value);
     assert_eq!(a.len(), 1);
     assert_eq!(a[key], value);
@@ -1809,9 +1810,9 @@ fn test_vacant_entry_key() {
     let key = "hello there";
     let value = "value goes here";
 
-    assert!(a.is_empty());
+    assert_eq!(a.height(), None);
     match a.entry(key) {
-        Occupied(_) => panic!(),
+        Occupied(_) => unreachable!(),
         Vacant(e) => {
             assert_eq!(key, *e.key());
             e.insert(value);
@@ -1822,6 +1823,36 @@ fn test_vacant_entry_key() {
     a.check();
 }
 
+#[test]
+fn test_vacant_entry_no_insert() {
+    let mut a = BTreeMap::<&str, ()>::new();
+    let key = "hello there";
+
+    // Non-allocated
+    assert_eq!(a.height(), None);
+    match a.entry(key) {
+        Occupied(_) => unreachable!(),
+        Vacant(e) => assert_eq!(key, *e.key()),
+    }
+    // Ensures the tree has no root.
+    assert_eq!(a.height(), None);
+    a.check();
+
+    // Allocated but still empty
+    a.insert(key, ());
+    a.remove(&key);
+    assert_eq!(a.height(), Some(0));
+    assert!(a.is_empty());
+    match a.entry(key) {
+        Occupied(_) => unreachable!(),
+        Vacant(e) => assert_eq!(key, *e.key()),
+    }
+    // Ensures the allocated root is not changed.
+    assert_eq!(a.height(), Some(0));
+    assert!(a.is_empty());
+    a.check();
+}
+
 #[test]
 fn test_first_last_entry() {
     let mut a = BTreeMap::new();

library/alloc/src/collections/btree/node.rs

@@ -213,7 +213,7 @@ unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Owned, K, V, Type>
 unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Dying, K, V, Type> {}
 
 impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
-    fn new_leaf() -> Self {
+    pub fn new_leaf() -> Self {
         Self::from_new_leaf(LeafNode::new())
     }
 
@@ -619,15 +619,16 @@ impl<K, V, Type> NodeRef<marker::Owned, K, V, Type> {
 }
 
 impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Leaf> {
-    /// Adds a key-value pair to the end of the node.
-    pub fn push(&mut self, key: K, val: V) {
+    /// Adds a key-value pair to the end of the node, and returns
+    /// the mutable reference of the inserted value.
+    pub fn push(&mut self, key: K, val: V) -> &mut V {
         let len = self.len_mut();
         let idx = usize::from(*len);
         assert!(idx < CAPACITY);
         *len += 1;
         unsafe {
             self.key_area_mut(idx).write(key);
-            self.val_area_mut(idx).write(val);
+            self.val_area_mut(idx).write(val)
         }
     }
 }