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

   1 use super::map::MIN_LEN;
   2 use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef};
   3
   4 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
   5     /// Removes a key-value pair from the tree, and returns that pair, as well as
   6     /// the leaf edge corresponding to that former pair. It's possible this empties
   7     /// a root node that is internal, which the caller should pop from the map
   8     /// holding the tree. The caller should also decrement the map's length.
   9     pub fn remove_kv_tracking<F: FnOnce()>(
  10         self,
  11         handle_emptied_internal_root: F,
  12     ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
  13         match self.force() {
  14             Leaf(node) => node.remove_leaf_kv(handle_emptied_internal_root),
  15             Internal(node) => node.remove_internal_kv(handle_emptied_internal_root),
  16         }
  17     }
  18 }
  19
  20 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::KV> {
  21     fn remove_leaf_kv<F: FnOnce()>(
  22         self,
  23         handle_emptied_internal_root: F,
  24     ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
  25         let (old_kv, mut pos) = self.remove();
  26         let len = pos.reborrow().into_node().len();
  27         if len < MIN_LEN {
  28             let idx = pos.idx();
  29             // We have to temporarily forget the child type, because there is no
  30             // distinct node type for the immediate parents of a leaf.
  31             let new_pos = match pos.into_node().forget_type().choose_parent_kv() {
  32                 Ok(Left(left_parent_kv)) => {
  33                     debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1);
  34                     if left_parent_kv.can_merge() {
  35                         left_parent_kv.merge_tracking_child_edge(Right(idx))
  36                     } else {
  37                         debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
  38                         left_parent_kv.steal_left(idx)
  39                     }
  40                 }
  41                 Ok(Right(right_parent_kv)) => {
  42                     debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1);
  43                     if right_parent_kv.can_merge() {
  44                         right_parent_kv.merge_tracking_child_edge(Left(idx))
  45                     } else {
  46                         debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
  47                         right_parent_kv.steal_right(idx)
  48                     }
  49                 }
  50                 Err(pos) => unsafe { Handle::new_edge(pos, idx) },
  51             };
  52             // SAFETY: `new_pos` is the leaf we started from or a sibling.
  53             pos = unsafe { new_pos.cast_to_leaf_unchecked() };
  54
  55             // Only if we merged, the parent (if any) has shrunk, but skipping
  56             // the following step otherwise does not pay off in benchmarks.
  57             //
  58             // SAFETY: We won't destroy or rearrange the leaf where `pos` is at
  59             // by handling its parent recursively; at worst we will destroy or
  60             // rearrange the parent through the grandparent, thus change the
  61             // link to the parent inside the leaf.
  62             if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() {
  63                 parent.into_node().handle_shrunk_node_recursively(handle_emptied_internal_root);
  64             }
  65         }
  66         (old_kv, pos)
  67     }
  68 }
  69
  70 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
  71     fn remove_internal_kv<F: FnOnce()>(
  72         self,
  73         handle_emptied_internal_root: F,
  74     ) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
  75         // Remove an adjacent KV from its leaf and then put it back in place of
  76         // the element we were asked to remove. Prefer the left adjacent KV,
  77         // for the reasons listed in `choose_parent_kv`.
  78         let left_leaf_kv = self.left_edge().descend().last_leaf_edge().left_kv();
  79         let left_leaf_kv = unsafe { left_leaf_kv.ok().unwrap_unchecked() };
  80         let (left_kv, left_hole) = left_leaf_kv.remove_leaf_kv(handle_emptied_internal_root);
  81
  82         // The internal node may have been stolen from or merged. Go back right
  83         // to find where the original KV ended up.
  84         let mut internal = unsafe { left_hole.next_kv().ok().unwrap_unchecked() };
  85         let old_kv = internal.replace_kv(left_kv.0, left_kv.1);
  86         let pos = internal.next_leaf_edge();
  87         (old_kv, pos)
  88     }
  89 }
  90
  91 impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
  92     /// Stocks up a possibly underfull internal node and its ancestors,
  93     /// until it reaches an ancestor that has elements to spare or is the root.
  94     fn handle_shrunk_node_recursively<F: FnOnce()>(mut self, handle_emptied_internal_root: F) {
  95         loop {
  96             self = match self.len() {
  97                 0 => {
  98                     // An empty node must be the root, because length is only
  99                     // reduced by one, and non-root underfull nodes are stocked up,
 100                     // so non-root nodes never have fewer than MIN_LEN - 1 elements.
 101                     debug_assert!(self.ascend().is_err());
 102                     handle_emptied_internal_root();
 103                     return;
 104                 }
 105                 1..MIN_LEN => {
 106                     if let Some(parent) = self.handle_underfull_node_locally() {
 107                         parent
 108                     } else {
 109                         return;
 110                     }
 111                 }
 112                 _ => return,
 113             }
 114         }
 115     }
 116
 117     /// Stocks up an underfull internal node, possibly at the cost of shrinking
 118     /// its parent instead, which is then returned.
 119     fn handle_underfull_node_locally(
 120         self,
 121     ) -> Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>> {
 122         match self.forget_type().choose_parent_kv() {
 123             Ok(Left(mut left_parent_kv)) => {
 124                 debug_assert_eq!(left_parent_kv.right_child_len(), MIN_LEN - 1);
 125                 if left_parent_kv.can_merge() {
 126                     let parent = left_parent_kv.merge_tracking_parent();
 127                     Some(parent)
 128                 } else {
 129                     debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
 130                     left_parent_kv.bulk_steal_left(1);
 131                     None
 132                 }
 133             }
 134             Ok(Right(mut right_parent_kv)) => {
 135                 debug_assert_eq!(right_parent_kv.left_child_len(), MIN_LEN - 1);
 136                 if right_parent_kv.can_merge() {
 137                     let parent = right_parent_kv.merge_tracking_parent();
 138                     Some(parent)
 139                 } else {
 140                     debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
 141                     right_parent_kv.bulk_steal_right(1);
 142                     None
 143                 }
 144             }
 145             Err(_) => None,
 146         }
 147     }
 148 }