2 files changed, 70 insertions, 19 deletions
diff --git a/library/alloc/src/collections/btree/map/entry.rs b/library/alloc/src/collections/btree/map/entry.rs
index 370b58864..e9366eec9 100644
--- a/library/alloc/src/collections/btree/map/entry.rs
+++ b/library/alloc/src/collections/btree/map/entry.rs
@@ -347,7 +347,7 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
     /// assert_eq!(map["poneyland"], 37);
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
-    pub fn insert(self, value: V) -> &'a mut V {
+    pub fn insert(mut self, value: V) -> &'a mut V {
         let out_ptr = match self.handle {
             None => {
                 // SAFETY: There is no tree yet so no reference to it exists.
@@ -358,25 +358,27 @@ impl<'a, K: Ord, V, A: Allocator + Clone> VacantEntry<'a, K, V, A> {
                 map.length = 1;
                 val_ptr
             }
-            Some(handle) => match handle.insert_recursing(self.key, value, self.alloc.clone()) {
-                (None, val_ptr) => {
-                    // SAFETY: We have consumed self.handle.
-                    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 dropped the
-                    // remaining reference to the tree, ins.left.
-                    let map = unsafe { self.dormant_map.awaken() };
-                    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);
-                    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
+            }
         };
+
         // 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 }
diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs
index 700b1463b..76c2f27b4 100644
--- a/library/alloc/src/collections/btree/map/tests.rs
+++ b/library/alloc/src/collections/btree/map/tests.rs
@@ -2336,3 +2336,52 @@ fn from_array() {
     let unordered_duplicates = BTreeMap::from([(3, 4), (1, 2), (1, 2)]);
     assert_eq!(map, unordered_duplicates);
 }
+
+#[test]
+fn test_cursor() {
+    let map = BTreeMap::from([(1, 'a'), (2, 'b'), (3, 'c')]);
+
+    let mut cur = map.lower_bound(Bound::Unbounded);
+    assert_eq!(cur.key(), Some(&1));
+    cur.move_next();
+    assert_eq!(cur.key(), Some(&2));
+    assert_eq!(cur.peek_next(), Some((&3, &'c')));
+    cur.move_prev();
+    assert_eq!(cur.key(), Some(&1));
+    assert_eq!(cur.peek_prev(), None);
+
+    let mut cur = map.upper_bound(Bound::Excluded(&1));
+    assert_eq!(cur.key(), None);
+    cur.move_next();
+    assert_eq!(cur.key(), Some(&1));
+    cur.move_prev();
+    assert_eq!(cur.key(), None);
+    assert_eq!(cur.peek_prev(), Some((&3, &'c')));
+}
+
+#[test]
+fn test_cursor_mut() {
+    let mut map = BTreeMap::from([(1, 'a'), (3, 'c'), (5, 'e')]);
+    let mut cur = map.lower_bound_mut(Bound::Excluded(&3));
+    assert_eq!(cur.key(), Some(&5));
+    cur.insert_before(4, 'd');
+    assert_eq!(cur.key(), Some(&5));
+    assert_eq!(cur.peek_prev(), Some((&4, &mut 'd')));
+    cur.move_next();
+    assert_eq!(cur.key(), None);
+    cur.insert_before(6, 'f');
+    assert_eq!(cur.key(), None);
+    assert_eq!(cur.remove_current(), None);
+    assert_eq!(cur.key(), None);
+    cur.insert_after(0, '?');
+    assert_eq!(cur.key(), None);
+    assert_eq!(map, BTreeMap::from([(0, '?'), (1, 'a'), (3, 'c'), (4, 'd'), (5, 'e'), (6, 'f')]));
+
+    let mut cur = map.upper_bound_mut(Bound::Included(&5));
+    assert_eq!(cur.key(), Some(&5));
+    assert_eq!(cur.remove_current(), Some((5, 'e')));
+    assert_eq!(cur.key(), Some(&6));
+    assert_eq!(cur.remove_current_and_move_back(), Some((6, 'f')));
+    assert_eq!(cur.key(), Some(&4));
+    assert_eq!(map, BTreeMap::from([(0, '?'), (1, 'a'), (3, 'c'), (4, 'd')]));
+}