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

   1 use super::node::{self, marker, ForceResult, Handle, NodeRef};
   2 use super::unwrap_unchecked;
   3 use core::mem;
   4 use core::ptr;
   5
   6 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
   7     /// Removes a key/value-pair from the map, and returns that pair, as well as
   8     /// the leaf edge corresponding to that former pair.
   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         let (old_kv, mut pos, was_internal) = match self.force() {
  14             ForceResult::Leaf(leaf) => {
  15                 let (old_kv, pos) = leaf.remove();
  16                 (old_kv, pos, false)
  17             }
  18             ForceResult::Internal(mut internal) => {
  19                 // Replace the location freed in the internal node with an
  20                 // adjacent KV, and remove that adjacent KV from its leaf.
  21                 // Always choose the adjacent KV on the left side because
  22                 // it is typically faster to pop an element from the end
  23                 // of the KV arrays without needing to shift other elements.
  24
  25                 let key_loc = internal.kv_mut().0 as *mut K;
  26                 let val_loc = internal.kv_mut().1 as *mut V;
  27
  28                 let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
  29                 let to_remove = unsafe { unwrap_unchecked(to_remove) };
  30
  31                 let (kv, pos) = to_remove.remove();
  32
  33                 let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) };
  34                 let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) };
  35
  36                 ((old_key, old_val), pos, true)
  37             }
  38         };
  39
  40         // Handle underflow
  41         let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() };
  42         let mut at_leaf = true;
  43         while cur_node.len() < node::MIN_LEN {
  44             match handle_underfull_node(cur_node) {
  45                 UnderflowResult::AtRoot => break,
  46                 UnderflowResult::Merged(edge, merged_with_left, offset) => {
  47                     // If we merged with our right sibling then our tracked
  48                     // position has not changed. However if we merged with our
  49                     // left sibling then our tracked position is now dangling.
  50                     if at_leaf && merged_with_left {
  51                         let idx = pos.idx() + offset;
  52                         let node = match unsafe { ptr::read(&edge).descend().force() } {
  53                             ForceResult::Leaf(leaf) => leaf,
  54                             ForceResult::Internal(_) => unreachable!(),
  55                         };
  56                         pos = unsafe { Handle::new_edge(node, idx) };
  57                     }
  58
  59                     let parent = edge.into_node();
  60                     if parent.len() == 0 {
  61                         // The parent that was just emptied must be the root,
  62                         // because nodes on a lower level would not have been
  63                         // left with a single child.
  64                         handle_emptied_internal_root();
  65                         break;
  66                     } else {
  67                         cur_node = parent.forget_type();
  68                         at_leaf = false;
  69                     }
  70                 }
  71                 UnderflowResult::Stole(stole_from_left) => {
  72                     // Adjust the tracked position if we stole from a left sibling
  73                     if stole_from_left && at_leaf {
  74                         // SAFETY: This is safe since we just added an element to our node.
  75                         unsafe {
  76                             pos.move_next_unchecked();
  77                         }
  78                     }
  79                     break;
  80                 }
  81             }
  82         }
  83
  84         // If we deleted from an internal node then we need to compensate for
  85         // the earlier swap and adjust the tracked position to point to the
  86         // next element.
  87         if was_internal {
  88             pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() };
  89         }
  90
  91         (old_kv, pos)
  92     }
  93 }
  94
  95 enum UnderflowResult<'a, K, V> {
  96     AtRoot,
  97     Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
  98     Stole(bool),
  99 }
 100
 101 fn handle_underfull_node<'a, K: 'a, V: 'a>(
 102     node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
 103 ) -> UnderflowResult<'_, K, V> {
 104     let parent = match node.ascend() {
 105         Ok(parent) => parent,
 106         Err(_) => return UnderflowResult::AtRoot,
 107     };
 108
 109     // Prefer the left KV if it exists. Merging with the left side is faster,
 110     // since merging happens towards the left and `node` has fewer elements.
 111     // Stealing from the left side is faster, since we can pop from the end of
 112     // the KV arrays.
 113     let (is_left, mut handle) = match parent.left_kv() {
 114         Ok(left) => (true, left),
 115         Err(parent) => {
 116             let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
 117             (false, right)
 118         }
 119     };
 120
 121     if handle.can_merge() {
 122         let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 };
 123         UnderflowResult::Merged(handle.merge(), is_left, offset)
 124     } else {
 125         if is_left {
 126             handle.steal_left();
 127         } else {
 128             handle.steal_right();
 129         }
 130         UnderflowResult::Stole(is_left)
 131     }
 132 }