Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix performance backlash of #74762 #75182

Closed
wants to merge 2 commits into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 15 additions & 55 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1255,11 +1255,7 @@ impl<K: Ord, V> BTreeMap<K, V> {
}
pub(super) fn drain_filter_inner(&mut self) -> DrainFilterInner<'_, K, V> {
let front = self.root.as_mut().map(|r| r.as_mut().first_leaf_edge());
DrainFilterInner {
length: &mut self.length,
cur_leaf_edge: front,
emptied_internal_root: false,
}
DrainFilterInner { length: &mut self.length, cur_leaf_edge: front }
}

/// Calculates the number of elements if it is incorrect.
Expand Down Expand Up @@ -1629,7 +1625,6 @@ where
pub(super) struct DrainFilterInner<'a, K: 'a, V: 'a> {
length: &'a mut usize,
cur_leaf_edge: Option<Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>>,
emptied_internal_root: bool,
}

#[unstable(feature = "btree_drain_filter", issue = "70530")]
Expand Down Expand Up @@ -1670,17 +1665,6 @@ where
}
}

impl<K, V> Drop for DrainFilterInner<'_, K, V> {
fn drop(&mut self) {
if self.emptied_internal_root {
if let Some(handle) = self.cur_leaf_edge.take() {
let root = handle.into_node().into_root_mut();
root.pop_internal_level();
}
}
}
}

impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
/// Allow Debug implementations to predict the next element.
pub(super) fn peek(&self) -> Option<(&K, &V)> {
Expand All @@ -1697,10 +1681,9 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> {
let (k, v) = kv.kv_mut();
if pred(k, v) {
*self.length -= 1;
let RemoveResult { old_kv, pos, emptied_internal_root } = kv.remove_kv_tracking();
self.cur_leaf_edge = Some(pos);
self.emptied_internal_root |= emptied_internal_root;
return Some(old_kv);
let (k, v, leaf_edge_location) = kv.remove_kv_tracking();
self.cur_leaf_edge = Some(leaf_edge_location);
return Some((k, v));
}
self.cur_leaf_edge = Some(kv.next_leaf_edge());
}
Expand Down Expand Up @@ -2668,31 +2651,17 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> {
fn remove_kv(self) -> (K, V) {
*self.length -= 1;

let RemoveResult { old_kv, pos, emptied_internal_root } = self.handle.remove_kv_tracking();
let root = pos.into_node().into_root_mut();
if emptied_internal_root {
root.pop_internal_level();
}
old_kv
let (old_key, old_val, _) = self.handle.remove_kv_tracking();
(old_key, old_val)
}
}

struct RemoveResult<'a, K, V> {
// Key and value removed.
old_kv: (K, V),
// Unique location at the leaf level that the removed KV lopgically collapsed into.
pos: Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>,
// Whether the remove left behind and empty internal root node, that should be removed
// using `pop_internal_level`.
emptied_internal_root: bool,
}

impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
/// Removes a key/value-pair from the tree, and returns that pair, as well as
/// the leaf edge corresponding to that former pair. It's possible this leaves
/// an empty internal root node, which the caller should subsequently pop from
/// the map holding the tree. The caller should also decrement the map's length.
fn remove_kv_tracking(self) -> RemoveResult<'a, K, V> {
/// Removes a key/value-pair from the map, and returns that pair, as well as
/// the leaf edge corresponding to that former pair.
fn remove_kv_tracking(
self,
) -> (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) => {
let (hole, old_key, old_val) = leaf.remove();
Expand Down Expand Up @@ -2721,7 +2690,6 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
};

// Handle underflow
let mut emptied_internal_root = false;
let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() };
let mut at_leaf = true;
while cur_node.len() < node::MIN_LEN {
Expand All @@ -2743,7 +2711,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
let parent = edge.into_node();
if parent.len() == 0 {
// This empty parent must be the root, and should be popped off the tree.
emptied_internal_root = true;
parent.into_root_mut().pop_internal_level();
break;
} else {
cur_node = parent.forget_type();
Expand All @@ -2770,7 +2738,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInter
pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() };
}

RemoveResult { old_kv: (old_key, old_val), pos, emptied_internal_root }
(old_key, old_val, pos)
}
}

Expand Down Expand Up @@ -2850,16 +2818,8 @@ fn handle_underfull_node<K, V>(
let (is_left, mut handle) = match parent.left_kv() {
Ok(left) => (true, left),
Err(parent) => {
match parent.right_kv() {
Ok(right) => (false, right),
Err(_) => {
// The underfull node has an empty parent, so it is the only child
// of an empty root. It is destined to become the new root, thus
// allowed to be underfull. The empty parent should be removed later
// by `pop_internal_level`.
return AtRoot;
}
}
let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
(false, right)
}
};

Expand Down