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

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