nordic-dev.net/rust
nordic-dev.net/rust
2
0
Fork 0
mirror of https://github.com/rust-lang/rust.git synced 2024-11-27 17:24:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

BTreeMap: split off most code of remove and split_off

Browse Source
This commit is contained in:
Stein Somers 2020-10-17 17:00:59 +02:00
parent c38ddb8040
commit 003516f91a
4 changed files with 239 additions and 222 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

223
library/alloc/src/collections/btree/map.rs
View File

@ -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::*;
@ -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();
@ -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);

2
library/alloc/src/collections/btree/mod.rs
View File

@ -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> {

132
library/alloc/src/collections/btree/remove.rs Normal file
View File

@ -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)
}
}

104
library/alloc/src/collections/btree/split.rs Normal file
View File

@ -0,0 +1,104 @@
use super::node::{self, ForceResult::*, Root};
use super::search::{self, SearchResult::*};
use core::borrow::Borrow;
impl<K, V> Root<K, V> {
pub fn split_off<Q: ?Sized + Ord>(&mut self, right_root: &mut Self, key: &Q)
where
K: Borrow<Q>,
{
debug_assert!(right_root.height() == 0);
debug_assert!(right_root.node_as_ref().len() == 0);
let left_root = self;
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();
}
/// Removes empty levels on the top, but keeps 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();
}
}