Skip to content

Commit

Permalink
Auto merge of #81362 - ssomers:btree_drainy_refactor_8, r=Mark-Simula…
Browse files Browse the repository at this point in the history
…crum

BTreeMap: gather and decompose reusable tree fixing functions

This is kind of pushing it as a standalone refactor, probably only useful for #81075 (or similar).
r? `@Mark-Simulacrum`
  • Loading branch information
bors committed Feb 22, 2021
2 parents 15598a8 + cbbdb44 commit a15f484
Show file tree
Hide file tree
Showing 5 changed files with 176 additions and 146 deletions.
23 changes: 1 addition & 22 deletions library/alloc/src/collections/btree/append.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use super::map::MIN_LEN;
use super::merge_iter::MergeIterInner;
use super::node::{self, ForceResult::*, Root};
use super::node::{self, Root};
use core::iter::FusedIterator;

impl<K, V> Root<K, V> {
Expand Down Expand Up @@ -83,26 +82,6 @@ impl<K, V> Root<K, V> {
}
self.fix_right_border_of_plentiful();
}

/// Stock up any underfull nodes on the right border of the tree.
/// The other nodes, those that are not the root nor a rightmost edge,
/// must have MIN_LEN elements to spare.
fn fix_right_border_of_plentiful(&mut self) {
let mut cur_node = self.borrow_mut();
while let Internal(internal) = cur_node.force() {
// Check if right-most child is underfull.
let mut last_kv = internal.last_kv().consider_for_balancing();
debug_assert!(last_kv.left_child_len() >= MIN_LEN * 2);
let right_child_len = last_kv.right_child_len();
if right_child_len < MIN_LEN {
// We need to steal.
last_kv.bulk_steal_left(MIN_LEN - right_child_len);
}

// Go further down.
cur_node = last_kv.into_right_child();
}
}
}

// An iterator for merging two sorted sequences into one
Expand Down
171 changes: 171 additions & 0 deletions library/alloc/src/collections/btree/fix.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,171 @@
use super::map::MIN_LEN;
use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef, Root};

impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
/// Stocks up a possibly underfull node by merging with or stealing from a
/// sibling. If succesful but at the cost of shrinking the parent node,
/// returns that shrunk parent node. Returns an `Err` if the node is
/// an empty root.
fn fix_node_through_parent(
self,
) -> Result<Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>>, Self> {
let len = self.len();
if len >= MIN_LEN {
Ok(None)
} else {
match self.choose_parent_kv() {
Ok(Left(mut left_parent_kv)) => {
if left_parent_kv.can_merge() {
let parent = left_parent_kv.merge_tracking_parent();
Ok(Some(parent))
} else {
left_parent_kv.bulk_steal_left(MIN_LEN - len);
Ok(None)
}
}
Ok(Right(mut right_parent_kv)) => {
if right_parent_kv.can_merge() {
let parent = right_parent_kv.merge_tracking_parent();
Ok(Some(parent))
} else {
right_parent_kv.bulk_steal_right(MIN_LEN - len);
Ok(None)
}
}
Err(root) => {
if len > 0 {
Ok(None)
} else {
Err(root)
}
}
}
}
}
}

impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
/// Stocks up a possibly underfull node, and if that causes its parent node
/// to shrink, stocks up the parent, recursively.
/// Returns `true` if it fixed the tree, `false` if it couldn't because the
/// root node became empty.
///
/// This method does not expect ancestors to already be underfull upon entry
/// and panics if it encounters an empty ancestor.
pub fn fix_node_and_affected_ancestors(mut self) -> bool {
loop {
match self.fix_node_through_parent() {
Ok(Some(parent)) => self = parent.forget_type(),
Ok(None) => return true,
Err(_) => return false,
}
}
}
}

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

/// Stocks up or merge away any underfull nodes on the right border of the
/// tree. The other nodes, those that are not the root nor a rightmost edge,
/// must already have at least MIN_LEN elements.
pub fn fix_right_border(&mut self) {
self.fix_top();
if self.len() > 0 {
self.borrow_mut().last_kv().fix_right_border_of_right_edge();
self.fix_top();
}
}

/// The symmetric clone of `fix_right_border`.
pub fn fix_left_border(&mut self) {
self.fix_top();
if self.len() > 0 {
self.borrow_mut().first_kv().fix_left_border_of_left_edge();
self.fix_top();
}
}

/// Stock up any underfull nodes on the right border of the tree.
/// The other nodes, those that are not the root nor a rightmost edge,
/// must be prepared to have up to MIN_LEN elements stolen.
pub fn fix_right_border_of_plentiful(&mut self) {
let mut cur_node = self.borrow_mut();
while let Internal(internal) = cur_node.force() {
// Check if right-most child is underfull.
let mut last_kv = internal.last_kv().consider_for_balancing();
debug_assert!(last_kv.left_child_len() >= MIN_LEN * 2);
let right_child_len = last_kv.right_child_len();
if right_child_len < MIN_LEN {
// We need to steal.
last_kv.bulk_steal_left(MIN_LEN - right_child_len);
}

// Go further down.
cur_node = last_kv.into_right_child();
}
}
}

impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
fn fix_left_border_of_left_edge(mut self) {
while let Internal(internal_kv) = self.force() {
self = internal_kv.fix_left_child().first_kv();
debug_assert!(self.reborrow().into_node().len() > MIN_LEN);
}
}

fn fix_right_border_of_right_edge(mut self) {
while let Internal(internal_kv) = self.force() {
self = internal_kv.fix_right_child().last_kv();
debug_assert!(self.reborrow().into_node().len() > MIN_LEN);
}
}
}

impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
/// Stocks up the left child, assuming the right child isn't underfull, and
/// provisions an extra element to allow merging its children in turn
/// without becoming underfull.
/// Returns the left child.
fn fix_left_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
let mut internal_kv = self.consider_for_balancing();
let left_len = internal_kv.left_child_len();
debug_assert!(internal_kv.right_child_len() >= MIN_LEN);
if internal_kv.can_merge() {
internal_kv.merge_tracking_child()
} else {
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
let count = (MIN_LEN + 1).saturating_sub(left_len);
if count > 0 {
internal_kv.bulk_steal_right(count);
}
internal_kv.into_left_child()
}
}

/// Stocks up the right child, assuming the left child isn't underfull, and
/// provisions an extra element to allow merging its children in turn
/// without becoming underfull.
/// Returns wherever the right child ended up.
fn fix_right_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
let mut internal_kv = self.consider_for_balancing();
let right_len = internal_kv.right_child_len();
debug_assert!(internal_kv.left_child_len() >= MIN_LEN);
if internal_kv.can_merge() {
internal_kv.merge_tracking_child()
} else {
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
let count = (MIN_LEN + 1).saturating_sub(right_len);
if count > 0 {
internal_kv.bulk_steal_left(count);
}
internal_kv.into_right_child()
}
}
}
1 change: 1 addition & 0 deletions library/alloc/src/collections/btree/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
mod append;
mod borrow;
mod fix;
pub mod map;
mod mem;
mod merge_iter;
Expand Down
63 changes: 3 additions & 60 deletions library/alloc/src/collections/btree/remove.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,9 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
// rearrange the parent through the grandparent, thus change the
// link to the parent inside the leaf.
if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() {
parent.into_node().handle_shrunk_node_recursively(handle_emptied_internal_root);
if !parent.into_node().forget_type().fix_node_and_affected_ancestors() {
handle_emptied_internal_root();
}
}
}
(old_kv, pos)
Expand All @@ -87,62 +89,3 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
(old_kv, pos)
}
}

impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
/// Stocks up a possibly underfull internal node and its ancestors,
/// until it reaches an ancestor that has elements to spare or is the root.
fn handle_shrunk_node_recursively<F: FnOnce()>(mut self, handle_emptied_internal_root: F) {
loop {
self = match self.len() {
0 => {
// An empty node must be the root, because length is only
// reduced by one, and non-root underfull nodes are stocked up,
// so non-root nodes never have fewer than MIN_LEN - 1 elements.
debug_assert!(self.ascend().is_err());
handle_emptied_internal_root();
return;
}
1..MIN_LEN => {
if let Some(parent) = self.handle_underfull_node_locally() {
parent
} else {
return;
}
}
_ => return,
}
}
}

/// Stocks up an underfull internal node, possibly at the cost of shrinking
/// its parent instead, which is then returned.
fn handle_underfull_node_locally(
self,
) -> Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>> {
match self.forget_type().choose_parent_kv() {
Ok(Left(mut left_parent_kv)) => {
debug_assert_eq!(left_parent_kv.right_child_len(), MIN_LEN - 1);
if left_parent_kv.can_merge() {
let parent = left_parent_kv.merge_tracking_parent();
Some(parent)
} else {
debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
left_parent_kv.bulk_steal_left(1);
None
}
}
Ok(Right(mut right_parent_kv)) => {
debug_assert_eq!(right_parent_kv.left_child_len(), MIN_LEN - 1);
if right_parent_kv.can_merge() {
let parent = right_parent_kv.merge_tracking_parent();
Some(parent)
} else {
debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
right_parent_kv.bulk_steal_right(1);
None
}
}
Err(_) => None,
}
}
}
64 changes: 0 additions & 64 deletions library/alloc/src/collections/btree/split.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
use super::map::MIN_LEN;
use super::node::{ForceResult::*, Root};
use super::search::SearchResult::*;
use core::borrow::Borrow;
Expand Down Expand Up @@ -70,67 +69,4 @@ impl<K, V> Root<K, V> {
}
root
}

/// 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.len() == 0 {
self.pop_internal_level();
}
}

/// Stock up or merge away any underfull nodes on the right border of the
/// tree. The other nodes, those that are not the root nor a rightmost edge,
/// must already have at least MIN_LEN elements.
fn fix_right_border(&mut self) {
self.fix_top();

{
let mut cur_node = self.borrow_mut();

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

if last_kv.can_merge() {
cur_node = last_kv.merge_tracking_child();
} else {
let right_len = last_kv.right_child_len();
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
if right_len < MIN_LEN + 1 {
last_kv.bulk_steal_left(MIN_LEN + 1 - right_len);
}
cur_node = last_kv.into_right_child();
}
debug_assert!(cur_node.len() > MIN_LEN);
}
}

self.fix_top();
}

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

{
let mut cur_node = self.borrow_mut();

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

if first_kv.can_merge() {
cur_node = first_kv.merge_tracking_child();
} else {
let left_len = first_kv.left_child_len();
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
if left_len < MIN_LEN + 1 {
first_kv.bulk_steal_right(MIN_LEN + 1 - left_len);
}
cur_node = first_kv.into_left_child();
}
debug_assert!(cur_node.len() > MIN_LEN);
}
}

self.fix_top();
}
}

0 comments on commit a15f484

Please sign in to comment.