BTreeMap: gather and decompose reusable tree fixing functions

[rust.git] / library / alloc / src / collections / btree / remove.rs

diff --git a/library/alloc/src/collections/btree/remove.rs b/library/alloc/src/collections/btree/remove.rs
index ff842197d19182782ae95c847e1d042abdd80061..6bc1252b9cb9d60604da628a2a5bf280f67e8579 100644 (file)
--- a/library/alloc/src/collections/btree/remove.rs
+++ b/library/alloc/src/collections/btree/remove.rs
@@ -1,6 +1,5 @@
 use super::map::MIN_LEN;
 use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef};
-use super::unwrap_unchecked;
 
 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
@@ -61,7 +60,9 @@ fn remove_leaf_kv<F: FnOnce()>(
             // rearrange the parent through the grandparent, thus change the
             // link to the parent inside the leaf.
             if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() {
-                parent.into_node().handle_shrunk_node_recursively(handle_emptied_internal_root);
+                if !parent.into_node().forget_type().fix_node_and_affected_ancestors() {
+                    handle_emptied_internal_root();
+                }
             }
         }
         (old_kv, pos)
@@ -77,73 +78,14 @@ fn remove_internal_kv<F: FnOnce()>(
         // the element we were asked to remove. Prefer the left adjacent KV,
         // for the reasons listed in `choose_parent_kv`.
         let left_leaf_kv = self.left_edge().descend().last_leaf_edge().left_kv();
-        let left_leaf_kv = unsafe { unwrap_unchecked(left_leaf_kv.ok()) };
+        let left_leaf_kv = unsafe { left_leaf_kv.ok().unwrap_unchecked() };
         let (left_kv, left_hole) = left_leaf_kv.remove_leaf_kv(handle_emptied_internal_root);
 
         // 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 mut internal = unsafe { left_hole.next_kv().ok().unwrap_unchecked() };
         let old_kv = internal.replace_kv(left_kv.0, left_kv.1);
         let pos = internal.next_leaf_edge();
         (old_kv, pos)
     }
 }
-
-impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
-    /// Stocks up a possibly underfull internal node and its ancestors,
-    /// until it reaches an ancestor that has elements to spare or is the root.
-    fn handle_shrunk_node_recursively<F: FnOnce()>(mut self, handle_emptied_internal_root: F) {
-        loop {
-            self = match self.len() {
-                0 => {
-                    // An empty node must be the root, because length is only
-                    // reduced by one, and non-root underfull nodes are stocked up,
-                    // so non-root nodes never have fewer than MIN_LEN - 1 elements.
-                    debug_assert!(self.ascend().is_err());
-                    handle_emptied_internal_root();
-                    return;
-                }
-                1..MIN_LEN => {
-                    if let Some(parent) = self.handle_underfull_node_locally() {
-                        parent
-                    } else {
-                        return;
-                    }
-                }
-                _ => return,
-            }
-        }
-    }
-
-    /// Stocks up an underfull internal node, possibly at the cost of shrinking
-    /// its parent instead, which is then returned.
-    fn handle_underfull_node_locally(
-        self,
-    ) -> Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>> {
-        match self.forget_type().choose_parent_kv() {
-            Ok(Left(mut left_parent_kv)) => {
-                debug_assert_eq!(left_parent_kv.right_child_len(), MIN_LEN - 1);
-                if left_parent_kv.can_merge() {
-                    let parent = left_parent_kv.merge_tracking_parent();
-                    Some(parent)
-                } else {
-                    debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
-                    left_parent_kv.bulk_steal_left(1);
-                    None
-                }
-            }
-            Ok(Right(mut right_parent_kv)) => {
-                debug_assert_eq!(right_parent_kv.left_child_len(), MIN_LEN - 1);
-                if right_parent_kv.can_merge() {
-                    let parent = right_parent_kv.merge_tracking_parent();
-                    Some(parent)
-                } else {
-                    debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
-                    right_parent_kv.bulk_steal_right(1);
-                    None
-                }
-            }
-            Err(_) => None,
-        }
-    }
-}