Rollup merge of #70795 - Amanieu:btree_remove_tracking, r=Mark-Simulacrum
Keep track of position when deleting from a BTreeMap This improves the performance of drain_filter and is needed for future Cursor support for BTreeMap. cc @ssomers r? @Mark-Simulacrum
This commit is contained in:
commit
2ec0e148f1
|
@ -1780,18 +1780,12 @@ where
|
||||||
where
|
where
|
||||||
F: FnMut(&K, &mut V) -> bool,
|
F: FnMut(&K, &mut V) -> bool,
|
||||||
{
|
{
|
||||||
while let Some(kv) = unsafe { self.next_kv() } {
|
while let Some(mut kv) = unsafe { self.next_kv() } {
|
||||||
let (k, v) = unsafe { ptr::read(&kv) }.into_kv_mut();
|
let (k, v) = kv.kv_mut();
|
||||||
if pred(k, v) {
|
if pred(k, v) {
|
||||||
*self.length -= 1;
|
*self.length -= 1;
|
||||||
let (k, v, leaf_edge_location) = kv.remove_kv_tracking();
|
let (k, v, leaf_edge_location) = kv.remove_kv_tracking();
|
||||||
// `remove_kv_tracking` has either preserved or invalidated `self.cur_leaf_edge`
|
self.cur_leaf_edge = Some(leaf_edge_location);
|
||||||
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),
|
|
||||||
}
|
|
||||||
};
|
|
||||||
return Some((k, v));
|
return Some((k, v));
|
||||||
}
|
}
|
||||||
self.cur_leaf_edge = Some(kv.next_leaf_edge());
|
self.cur_leaf_edge = Some(kv.next_leaf_edge());
|
||||||
|
@ -2698,28 +2692,26 @@ 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> {
|
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
|
/// 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:
|
/// 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.
|
|
||||||
fn remove_kv_tracking(
|
fn remove_kv_tracking(
|
||||||
self,
|
self,
|
||||||
) -> (K, V, Option<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>>) {
|
) -> (K, V, Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
|
||||||
let mut levels_down_handled: isize;
|
let (mut pos, old_key, old_val, was_internal) = match self.force() {
|
||||||
let (small_leaf, old_key, old_val) = match self.force() {
|
|
||||||
Leaf(leaf) => {
|
Leaf(leaf) => {
|
||||||
levels_down_handled = 1; // handled at same level, but affects only the right side
|
|
||||||
let (hole, old_key, old_val) = leaf.remove();
|
let (hole, old_key, old_val) = leaf.remove();
|
||||||
(hole.into_node(), old_key, old_val)
|
(hole, old_key, old_val, false)
|
||||||
}
|
}
|
||||||
Internal(mut internal) => {
|
Internal(mut internal) => {
|
||||||
// Replace the location freed in the internal node with the next KV,
|
// Replace the location freed in the internal node with the next KV,
|
||||||
// and remove that next KV from its leaf.
|
// 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 key_loc = internal.kv_mut().0 as *mut K;
|
||||||
let val_loc = internal.kv_mut().1 as *mut V;
|
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 to_remove = unsafe { unwrap_unchecked(to_remove) };
|
||||||
|
|
||||||
let (hole, key, val) = to_remove.remove();
|
let (hole, key, val) = to_remove.remove();
|
||||||
|
@ -2727,50 +2719,69 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
|
||||||
let old_key = unsafe { mem::replace(&mut *key_loc, key) };
|
let old_key = unsafe { mem::replace(&mut *key_loc, key) };
|
||||||
let old_val = unsafe { mem::replace(&mut *val_loc, val) };
|
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
|
// 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 {
|
while cur_node.len() < node::MIN_LEN {
|
||||||
match handle_underfull_node(cur_node) {
|
match handle_underfull_node(cur_node) {
|
||||||
AtRoot(root) => {
|
AtRoot => break,
|
||||||
cur_node = root;
|
Merged(edge, merged_with_left, offset) => {
|
||||||
break;
|
// If we merged with our right sibling then our tracked
|
||||||
}
|
// position has not changed. However if we merged with our
|
||||||
EmptyParent(_) => unreachable!(),
|
// left sibling then our tracked position is now dangling.
|
||||||
Merged(parent) => {
|
if at_leaf && merged_with_left {
|
||||||
levels_down_handled -= 1;
|
let idx = pos.idx() + offset;
|
||||||
|
let node = match unsafe { ptr::read(&edge).descend().force() } {
|
||||||
|
Leaf(leaf) => leaf,
|
||||||
|
Internal(_) => unreachable!(),
|
||||||
|
};
|
||||||
|
pos = unsafe { Handle::new_edge(node, idx) };
|
||||||
|
}
|
||||||
|
|
||||||
|
let parent = edge.into_node();
|
||||||
if parent.len() == 0 {
|
if parent.len() == 0 {
|
||||||
// We must be at the root
|
// We must be at the root
|
||||||
let root = parent.into_root_mut();
|
parent.into_root_mut().pop_level();
|
||||||
root.pop_level();
|
|
||||||
cur_node = root.as_mut();
|
|
||||||
break;
|
break;
|
||||||
} else {
|
} else {
|
||||||
cur_node = parent.forget_type();
|
cur_node = parent.forget_type();
|
||||||
|
at_leaf = false;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Stole(internal_node) => {
|
Stole(stole_from_left) => {
|
||||||
levels_down_handled -= 1;
|
// Adjust the tracked position if we stole from a left sibling
|
||||||
cur_node = internal_node.forget_type();
|
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.
|
// This internal node might be underfull, but only if it's the root.
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let leaf_edge_location = if levels_down_handled > 0 { None } else { Some(cur_node) };
|
// If we deleted from an internal node then we need to compensate for
|
||||||
(old_key, old_val, leaf_edge_location)
|
// 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> {
|
enum UnderflowResult<'a, K, V> {
|
||||||
AtRoot(NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>),
|
AtRoot,
|
||||||
EmptyParent(NodeRef<marker::Mut<'a>, K, V, marker::Internal>),
|
Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
|
||||||
Merged(NodeRef<marker::Mut<'a>, K, V, marker::Internal>),
|
Stole(bool),
|
||||||
Stole(NodeRef<marker::Mut<'a>, K, V, marker::Internal>),
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn handle_underfull_node<K, V>(
|
fn handle_underfull_node<K, V>(
|
||||||
|
@ -2778,28 +2789,27 @@ fn handle_underfull_node<K, V>(
|
||||||
) -> UnderflowResult<'_, K, V> {
|
) -> UnderflowResult<'_, K, V> {
|
||||||
let parent = match node.ascend() {
|
let parent = match node.ascend() {
|
||||||
Ok(parent) => parent,
|
Ok(parent) => parent,
|
||||||
Err(root) => return AtRoot(root),
|
Err(_) => return AtRoot,
|
||||||
};
|
};
|
||||||
|
|
||||||
let (is_left, mut handle) = match parent.left_kv() {
|
let (is_left, mut handle) = match parent.left_kv() {
|
||||||
Ok(left) => (true, left),
|
Ok(left) => (true, left),
|
||||||
Err(parent) => match parent.right_kv() {
|
Err(parent) => {
|
||||||
Ok(right) => (false, right),
|
let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
|
||||||
Err(parent) => {
|
(false, right)
|
||||||
return EmptyParent(parent.into_node());
|
}
|
||||||
}
|
|
||||||
},
|
|
||||||
};
|
};
|
||||||
|
|
||||||
if handle.can_merge() {
|
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 {
|
} else {
|
||||||
if is_left {
|
if is_left {
|
||||||
handle.steal_left();
|
handle.steal_left();
|
||||||
} else {
|
} else {
|
||||||
handle.steal_right();
|
handle.steal_right();
|
||||||
}
|
}
|
||||||
Stole(handle.into_node())
|
Stole(is_left)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -723,6 +723,11 @@ impl<Node, Type> Handle<Node, Type> {
|
||||||
pub fn into_node(self) -> Node {
|
pub fn into_node(self) -> 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> {
|
impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV> {
|
||||||
|
|
Loading…
Reference in New Issue