BTreeMap: capture a recurring use pattern as replace_kv

2 files changed, 10 insertions, 8 deletions

diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs
index 31809fde57b..ae5831d5140 100644
--- a/library/alloc/src/collections/btree/node.rs
+++ b/library/alloc/src/collections/btree/node.rs
@@ -1159,6 +1159,12 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>
             (key, val)
         }
     }
+
+    /// Replace the key and value that the KV handle refers to.
+    pub fn replace_kv(&mut self, k: K, v: V) -> (K, V) {
+        let (key, val) = self.kv_mut();
+        (mem::replace(key, k), mem::replace(val, v))
+    }
 }
 
 impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
@@ -1432,8 +1438,7 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
         unsafe {
             let (k, v, edge) = self.left_child.pop();
 
-            let k = mem::replace(self.parent.kv_mut().0, k);
-            let v = mem::replace(self.parent.kv_mut().1, v);
+            let (k, v) = self.parent.replace_kv(k, v);
 
             match self.right_child.reborrow_mut().force() {
                 ForceResult::Leaf(mut leaf) => leaf.push_front(k, v),
@@ -1455,8 +1460,7 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
         unsafe {
             let (k, v, edge) = self.right_child.pop_front();
 
-            let k = mem::replace(self.parent.kv_mut().0, k);
-            let v = mem::replace(self.parent.kv_mut().1, v);
+            let (k, v) = self.parent.replace_kv(k, v);
 
             match self.left_child.reborrow_mut().force() {
                 ForceResult::Leaf(mut leaf) => leaf.push(k, v),
diff --git a/library/alloc/src/collections/btree/remove.rs b/library/alloc/src/collections/btree/remove.rs
index 5ae1c1fcab8..1da0a060806 100644
--- a/library/alloc/src/collections/btree/remove.rs
+++ b/library/alloc/src/collections/btree/remove.rs
@@ -1,7 +1,6 @@
 use super::map::MIN_LEN;
 use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef};
 use super::unwrap_unchecked;
-use core::mem;
 
 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
     /// Removes a key-value pair from the tree, and returns that pair, as well as
@@ -84,10 +83,9 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
         // The internal node may have been stolen from or merged. Go back right
         // to find where the original KV ended up.
         let mut internal = unsafe { unwrap_unchecked(left_hole.next_kv().ok()) };
-        let old_key = mem::replace(internal.kv_mut().0, left_kv.0);
-        let old_val = mem::replace(internal.kv_mut().1, left_kv.1);
+        let old_kv = internal.replace_kv(left_kv.0, left_kv.1);
         let pos = internal.next_leaf_edge();
-        ((old_key, old_val), pos)
+        (old_kv, pos)
     }
 }