Keep track of position when deleting from a BTreeMap

src/liballoc/collections/btree/map.rs

@@ -1780,18 +1780,12 @@ where
  where
  F: FnMut(&K, &mut V) -> bool,
  {
- while let Some(kv) = unsafe { self.next_kv() } {
- let (k, v) = unsafe { ptr::read(&kv) }.into_kv_mut();
+ while let Some(mut kv) = unsafe { self.next_kv() } {
+ let (k, v) = kv.kv_mut();
  if pred(k, v) {
  *self.length -= 1;
  let (k, v, leaf_edge_location) = kv.remove_kv_tracking();
- // `remove_kv_tracking` has either preserved or invalidated `self.cur_leaf_edge`
- if let Some(node) = leaf_edge_location {
- match search::search_tree(node, &k) {
- search::SearchResult::Found(_) => unreachable!(),
- search::SearchResult::GoDown(leaf) => self.cur_leaf_edge = Some(leaf),
- }
- };
+ self.cur_leaf_edge = Some(leaf_edge_location);
  return Some((k, v));
  }
  self.cur_leaf_edge = Some(kv.next_leaf_edge());
@@ -2698,79 +2692,99 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
 
 impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
  /// Removes a key/value-pair from the map, and returns that pair, as well as
- /// the whereabouts of the leaf edge corresponding to that former pair:
- /// if None is returned, the leaf edge is still the left leaf edge of the KV handle;
- /// if a node is returned, it heads the subtree where the leaf edge may be found.
+ /// the leaf edge corresponding to that former pair.
  fn remove_kv_tracking(
  self,
- ) -> (K, V, Option<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>>) {
- let mut levels_down_handled: isize;
- let (small_leaf, old_key, old_val) = match self.force() {
+ ) -> (K, V, Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
+ let (mut pos, old_key, old_val, was_internal) = match self.force() {
  Leaf(leaf) => {
- levels_down_handled = 1; // handled at same level, but affects only the right side
  let (hole, old_key, old_val) = leaf.remove();
- (hole.into_node(), old_key, old_val)
+ (hole, old_key, old_val, false)
  }
  Internal(mut internal) => {
  // Replace the location freed in the internal node with the next KV,
  // and remove that next KV from its leaf.
- levels_down_handled = unsafe { ptr::read(&internal).into_node().height() } as isize;
 
  let key_loc = internal.kv_mut().0 as *mut K;
  let val_loc = internal.kv_mut().1 as *mut V;
 
- let to_remove = internal.right_edge().descend().first_leaf_edge().right_kv().ok();
+ // Deleting from the left side is typically faster since we can
+ // just pop an element from the end of the KV array without
+ // needing to shift the other values.
+ let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
  let to_remove = unsafe { unwrap_unchecked(to_remove) };
 
  let (hole, key, val) = to_remove.remove();
 
  let old_key = unsafe { mem::replace(&mut *key_loc, key) };
  let old_val = unsafe { mem::replace(&mut *val_loc, val) };
 
- (hole.into_node(), old_key, old_val)
+ (hole, old_key, old_val, true)
  }
  };
 
  // Handle underflow
- let mut cur_node = small_leaf.forget_type();
+ let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() };
+ let mut at_leaf = true;
  while cur_node.len() < node::MIN_LEN {
  match handle_underfull_node(cur_node) {
- AtRoot(root) => {
- cur_node = root;
- break;
- }
+ AtRoot(_) => break,
  EmptyParent(_) => unreachable!(),
- Merged(parent) => {
- levels_down_handled -= 1;
+ Merged(edge, merged_with_left, offset) => {
+ // If we merged with our right sibling then our tracked
+ // position has not changed. However if we merged with our
+ // left sibling then our tracked position is now dangling.
+ if at_leaf && merged_with_left {
+ let idx = pos.idx() + offset;
+ let node = match unsafe { ptr::read(&edge).descend().force() } {
+ Leaf(leaf) => leaf,
+ Internal(_) => unreachable!(),
+ };
+ debug_assert!(idx <= node.len());
+ pos = unsafe { Handle::new_edge(node, idx) };
+ }
+
+ let parent = edge.into_node();
  if parent.len() == 0 {
  // We must be at the root
- let root = parent.into_root_mut();
- root.pop_level();
- cur_node = root.as_mut();
+ parent.into_root_mut().pop_level();
  break;
  } else {
  cur_node = parent.forget_type();
+ at_leaf = false;
  }
  }
- Stole(internal_node) => {
- levels_down_handled -= 1;
- cur_node = internal_node.forget_type();
+ Stole(_, stole_from_left) => {
+ // Adjust the tracked position if we stole from a left sibling
+ if stole_from_left && at_leaf {
+ // SAFETY: This is safe since we just added an element to our node.
+ unsafe {
+ pos.next_unchecked();
+ }
+ }
+
  // This internal node might be underfull, but only if it's the root.
  break;
  }
  }
  }
 
- let leaf_edge_location = if levels_down_handled > 0 { None } else { Some(cur_node) };
- (old_key, old_val, leaf_edge_location)
+ // If we deleted from an internal node then we need to compensate for
+ // the earlier swap and adjust the tracked position to point to the
+ // next element.
+ if was_internal {
+ pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() };
+ }
+
+ (old_key, old_val, pos)
  }
 }
 
 enum UnderflowResult<'a, K, V> {
  AtRoot(NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>),
  EmptyParent(NodeRef<marker::Mut<'a>, K, V, marker::Internal>),
- Merged(NodeRef<marker::Mut<'a>, K, V, marker::Internal>),
- Stole(NodeRef<marker::Mut<'a>, K, V, marker::Internal>),
+ Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
+ Stole(NodeRef<marker::Mut<'a>, K, V, marker::Internal>, bool),
 }
 
 fn handle_underfull_node<K, V>(
@@ -2792,14 +2806,15 @@ fn handle_underfull_node<K, V>(
  };
 
  if handle.can_merge() {
- Merged(handle.merge().into_node())
+ let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 };
+ Merged(handle.merge(), is_left, offset)
  } else {
  if is_left {
  handle.steal_left();
  } else {
  handle.steal_right();
  }
- Stole(handle.into_node())
+ Stole(handle.into_node(), is_left)
  }
 }
 

src/liballoc/collections/btree/node.rs

@@ -723,6 +723,11 @@ impl<Node, Type> Handle<Node, Type> {
  pub fn into_node(self) -> Node {
  self.node
  }
+
+ /// Returns the position of this handle in the node.
+ pub fn idx(&self) -> usize {
+ self.idx
+ }
 }
 
 impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV> {