From 4f4c8ba0ae50878d0287424eb107fce9d3d8ecd6 Mon Sep 17 00:00:00 2001
From: Stein Somers <git@steinsomers.be>
Date: Fri, 3 Jan 2020 18:38:59 +0100
Subject: [PATCH] divert get_mut to operate on into_key_slice_mut

---
 src/liballoc/collections/btree/map.rs    |  2 +-
 src/liballoc/collections/btree/node.rs   |  2 +-
 src/liballoc/collections/btree/search.rs | 54 ++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index fa8aae04011ed..784f4954a89b5 100644
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -743,7 +743,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
         K: Borrow<Q>,
         Q: Ord,
     {
-        match search::search_tree(self.root.as_mut(), key) {
+        match search::search_tree_mut(self.root.as_mut(), key) {
             Found(handle) => Some(handle.into_kv_mut().1),
             GoDown(_) => None,
         }
diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs
index 0a2849f6e3961..e43d893ce4531 100644
--- a/src/liballoc/collections/btree/node.rs
+++ b/src/liballoc/collections/btree/node.rs
@@ -497,7 +497,7 @@ impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, K, V, Type> {
         self.node.as_ptr()
     }
 
-    fn keys_mut(&mut self) -> &mut [K] {
+    pub fn keys_mut(&mut self) -> &mut [K] {
         unsafe { self.reborrow_mut().into_key_slice_mut() }
     }
 
diff --git a/src/liballoc/collections/btree/search.rs b/src/liballoc/collections/btree/search.rs
index 3f3c49a2ef875..6d9611eac56e6 100644
--- a/src/liballoc/collections/btree/search.rs
+++ b/src/liballoc/collections/btree/search.rs
@@ -32,6 +32,28 @@ where
     }
 }
 
+pub fn search_tree_mut<'a, K, V, Q: ?Sized>(
+    mut node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
+    key: &Q,
+) -> SearchResult<marker::Mut<'a>, K, V, marker::LeafOrInternal, marker::Leaf>
+where
+    Q: Ord,
+    K: Borrow<Q>,
+{
+    loop {
+        match search_node_mut(node, key) {
+            Found(handle) => return Found(handle),
+            GoDown(handle) => match handle.force() {
+                Leaf(leaf) => return GoDown(leaf),
+                Internal(internal) => {
+                    node = internal.descend();
+                    continue;
+                }
+            },
+        }
+    }
+}
+
 pub fn search_node<BorrowType, K, V, Type, Q: ?Sized>(
     node: NodeRef<BorrowType, K, V, Type>,
     key: &Q,
@@ -46,6 +68,20 @@ where
     }
 }
 
+pub fn search_node_mut<'a, K, V, Type, Q: ?Sized>(
+    mut node: NodeRef<marker::Mut<'a>, K, V, Type>,
+    key: &Q,
+) -> SearchResult<marker::Mut<'a>, K, V, Type, Type>
+where
+    Q: Ord,
+    K: Borrow<Q>,
+{
+    match search_linear_mut(&mut node, key) {
+        (idx, true) => Found(Handle::new_kv(node, idx)),
+        (idx, false) => SearchResult::GoDown(Handle::new_edge(node, idx)),
+    }
+}
+
 pub fn search_linear<BorrowType, K, V, Type, Q: ?Sized>(
     node: &NodeRef<BorrowType, K, V, Type>,
     key: &Q,
@@ -63,3 +99,21 @@ where
     }
     (node.keys().len(), false)
 }
+
+pub fn search_linear_mut<'a, K, V, Type, Q: ?Sized>(
+    node: &mut NodeRef<marker::Mut<'a>, K, V, Type>,
+    key: &Q,
+) -> (usize, bool)
+where
+    Q: Ord,
+    K: Borrow<Q>,
+{
+    for (i, k) in node.keys_mut().iter().enumerate() {
+        match key.cmp(k.borrow()) {
+            Ordering::Greater => {}
+            Ordering::Equal => return (i, true),
+            Ordering::Less => return (i, false),
+        }
+    }
+    (node.keys().len(), false)
+}