diff --git a/library/alloc/src/collections/btree/append.rs b/library/alloc/src/collections/btree/append.rs index 1d6488dd2dfe3..a30a21db5347f 100644 --- a/library/alloc/src/collections/btree/append.rs +++ b/library/alloc/src/collections/btree/append.rs @@ -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 Root { @@ -83,26 +82,6 @@ impl Root { } 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 diff --git a/library/alloc/src/collections/btree/fix.rs b/library/alloc/src/collections/btree/fix.rs new file mode 100644 index 0000000000000..af87a9b956a17 --- /dev/null +++ b/library/alloc/src/collections/btree/fix.rs @@ -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, 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, 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, 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 Root { + /// 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, 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, 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, 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, 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() + } + } +} diff --git a/library/alloc/src/collections/btree/mod.rs b/library/alloc/src/collections/btree/mod.rs index 421f842dab0ac..f74172c7d976f 100644 --- a/library/alloc/src/collections/btree/mod.rs +++ b/library/alloc/src/collections/btree/mod.rs @@ -1,5 +1,6 @@ mod append; mod borrow; +mod fix; pub mod map; mod mem; mod merge_iter; diff --git a/library/alloc/src/collections/btree/remove.rs b/library/alloc/src/collections/btree/remove.rs index 9cd016fa62f0e..6bc1252b9cb9d 100644 --- a/library/alloc/src/collections/btree/remove.rs +++ b/library/alloc/src/collections/btree/remove.rs @@ -60,7 +60,9 @@ impl<'a, K: 'a, V: 'a> Handle, 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) @@ -87,62 +89,3 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::Internal>, (old_kv, pos) } } - -impl<'a, K: 'a, V: 'a> NodeRef, 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(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, 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, - } - } -} diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs index 1921982464ae4..bec495a72a4a8 100644 --- a/library/alloc/src/collections/btree/split.rs +++ b/library/alloc/src/collections/btree/split.rs @@ -1,4 +1,3 @@ -use super::map::MIN_LEN; use super::node::{ForceResult::*, Root}; use super::search::SearchResult::*; use core::borrow::Borrow; @@ -70,67 +69,4 @@ impl Root { } 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(); - } }