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

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