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.

Sign up for GitHub

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

Jump to bottom

BTreeMap: split off most code of remove and split_off #78056

Merged
merged 1 commit into from
Oct 21, 2020
Merged

BTreeMap: split off most code of remove and split_off #78056

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
223 changes: 1 addition & 222 deletions library/alloc/src/collections/btree/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,6 @@ use super::node::{self, marker, ForceResult::*, Handle, NodeRef};
use super::search::{self, SearchResult::*};
use super::unwrap_unchecked;

use UnderflowResult::*;

mod entry;
pub use entry::{Entry, OccupiedEntry, VacantEntry};
use Entry::*;
Expand Down Expand Up @@ -1154,40 +1152,8 @@ impl<K: Ord, V> BTreeMap<K, V> {

let mut right = Self::new();
let right_root = Self::ensure_is_owned(&mut right.root);
for _ in 0..left_root.height() {
right_root.push_internal_level();
}

{
let mut left_node = left_root.node_as_mut();
let mut right_node = right_root.node_as_mut();

loop {
let mut split_edge = match search::search_node(left_node, key) {
// key is going to the right tree
Found(handle) => handle.left_edge(),
GoDown(handle) => handle,
};

split_edge.move_suffix(&mut right_node);

match (split_edge.force(), right_node.force()) {
(Internal(edge), Internal(node)) => {
left_node = edge.descend();
right_node = node.first_edge().descend();
}
(Leaf(_), Leaf(_)) => {
break;
}
_ => {
unreachable!();
}
}
}
}

left_root.fix_right_border();
right_root.fix_left_border();
left_root.split_off(right_root, key);

if left_root.height() < right_root.height() {
self.length = left_root.node_as_ref().calc_length();
Expand Down Expand Up @@ -2250,193 +2216,6 @@ impl<K, V> BTreeMap<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 leaf edge corresponding to that former pair.
fn remove_kv_tracking<F: FnOnce()>(
self,
handle_emptied_internal_root: F,
) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
let (old_kv, mut pos, was_internal) = match self.force() {
Leaf(leaf) => {
let (old_kv, pos) = leaf.remove();
(old_kv, pos, false)
}
Internal(mut internal) => {
// Replace the location freed in the internal node with an
// adjacent KV, and remove that adjacent KV from its leaf.
// Always choose the adjacent KV on the left side because
// it is typically faster to pop an element from the end
// of the KV arrays without needing to shift other elements.

let key_loc = internal.kv_mut().0 as *mut K;
let val_loc = internal.kv_mut().1 as *mut V;

let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
let to_remove = unsafe { unwrap_unchecked(to_remove) };

let (kv, pos) = to_remove.remove();

let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) };
let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) };

((old_key, old_val), pos, true)
}
};

// Handle underflow
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 => break,
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!(),
};
pos = unsafe { Handle::new_edge(node, idx) };
}

let parent = edge.into_node();
if parent.len() == 0 {
// The parent that was just emptied must be the root,
// because nodes on a lower level would not have been
// left with a single child.
handle_emptied_internal_root();
break;
} else {
cur_node = parent.forget_type();
at_leaf = false;
}
}
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.move_next_unchecked();
}
}
break;
}
}
}

// 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_kv, pos)
}
}

impl<K, V> node::Root<K, V> {
/// Removes empty levels on the top, but keep an empty leaf if the entire tree is empty.
fn fix_top(&mut self) {
while self.height() > 0 && self.node_as_ref().len() == 0 {
self.pop_internal_level();
}
}

fn fix_right_border(&mut self) {
self.fix_top();

{
let mut cur_node = self.node_as_mut();

while let Internal(node) = cur_node.force() {
let mut last_kv = node.last_kv();

if last_kv.can_merge() {
cur_node = last_kv.merge().descend();
} else {
let right_len = last_kv.reborrow().right_edge().descend().len();
// `MINLEN + 1` to avoid readjust if merge happens on the next level.
if right_len < node::MIN_LEN + 1 {
last_kv.bulk_steal_left(node::MIN_LEN + 1 - right_len);
}
cur_node = last_kv.right_edge().descend();
}
}
}

self.fix_top();
}

/// The symmetric clone of `fix_right_border`.
fn fix_left_border(&mut self) {
self.fix_top();

{
let mut cur_node = self.node_as_mut();

while let Internal(node) = cur_node.force() {
let mut first_kv = node.first_kv();

if first_kv.can_merge() {
cur_node = first_kv.merge().descend();
} else {
let left_len = first_kv.reborrow().left_edge().descend().len();
if left_len < node::MIN_LEN + 1 {
first_kv.bulk_steal_right(node::MIN_LEN + 1 - left_len);
}
cur_node = first_kv.left_edge().descend();
}
}
}

self.fix_top();
}
}

enum UnderflowResult<'a, K, V> {
AtRoot,
Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
Stole(bool),
}

fn handle_underfull_node<'a, K: 'a, V: 'a>(
node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
) -> UnderflowResult<'_, K, V> {
let parent = match node.ascend() {
Ok(parent) => parent,
Err(_) => return AtRoot,
};

// Prefer the left KV if it exists. Merging with the left side is faster,
// since merging happens towards the left and `node` has fewer elements.
// Stealing from the left side is faster, since we can pop from the end of
// the KV arrays.
let (is_left, mut handle) = match parent.left_kv() {
Ok(left) => (true, left),
Err(parent) => {
let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
(false, right)
}
};

if handle.can_merge() {
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(is_left)
}
}

impl<K: Ord, V, I: Iterator<Item = (K, V)>> Iterator for MergeIter<K, V, I> {
type Item = (K, V);

Expand Down
2 changes: 2 additions & 0 deletions library/alloc/src/collections/btree/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@ mod borrow;
pub mod map;
mod navigate;
mod node;
mod remove;
mod search;
pub mod set;
mod split;

#[doc(hidden)]
trait Recover<Q: ?Sized> {
Expand Down
132 changes: 132 additions & 0 deletions library/alloc/src/collections/btree/remove.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
use super::node::{self, marker, ForceResult, Handle, NodeRef};
use super::unwrap_unchecked;
use core::mem;
use core::ptr;

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 leaf edge corresponding to that former pair.
pub fn remove_kv_tracking<F: FnOnce()>(
self,
handle_emptied_internal_root: F,
) -> ((K, V), Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge>) {
let (old_kv, mut pos, was_internal) = match self.force() {
ForceResult::Leaf(leaf) => {
let (old_kv, pos) = leaf.remove();
(old_kv, pos, false)
}
ForceResult::Internal(mut internal) => {
// Replace the location freed in the internal node with an
// adjacent KV, and remove that adjacent KV from its leaf.
// Always choose the adjacent KV on the left side because
// it is typically faster to pop an element from the end
// of the KV arrays without needing to shift other elements.

let key_loc = internal.kv_mut().0 as *mut K;
let val_loc = internal.kv_mut().1 as *mut V;

let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok();
let to_remove = unsafe { unwrap_unchecked(to_remove) };

let (kv, pos) = to_remove.remove();

let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) };
let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) };

((old_key, old_val), pos, true)
}
};

// Handle underflow
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) {
UnderflowResult::AtRoot => break,
UnderflowResult::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() } {
ForceResult::Leaf(leaf) => leaf,
ForceResult::Internal(_) => unreachable!(),
};
pos = unsafe { Handle::new_edge(node, idx) };
}

let parent = edge.into_node();
if parent.len() == 0 {
// The parent that was just emptied must be the root,
// because nodes on a lower level would not have been
// left with a single child.
handle_emptied_internal_root();
break;
} else {
cur_node = parent.forget_type();
at_leaf = false;
}
}
UnderflowResult::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.move_next_unchecked();
}
}
break;
}
}
}

// 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_kv, pos)
}
}

enum UnderflowResult<'a, K, V> {
AtRoot,
Merged(Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::Edge>, bool, usize),
Stole(bool),
}

fn handle_underfull_node<'a, K: 'a, V: 'a>(
node: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
) -> UnderflowResult<'_, K, V> {
let parent = match node.ascend() {
Ok(parent) => parent,
Err(_) => return UnderflowResult::AtRoot,
};

// Prefer the left KV if it exists. Merging with the left side is faster,
// since merging happens towards the left and `node` has fewer elements.
// Stealing from the left side is faster, since we can pop from the end of
// the KV arrays.
let (is_left, mut handle) = match parent.left_kv() {
Ok(left) => (true, left),
Err(parent) => {
let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) };
(false, right)
}
};

if handle.can_merge() {
let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 };
UnderflowResult::Merged(handle.merge(), is_left, offset)
} else {
if is_left {
handle.steal_left();
} else {
handle.steal_right();
}
UnderflowResult::Stole(is_left)
}
}
Loading