From ff3cab4c3f69274000cdc3e3958810d99f4abb99 Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Wed, 18 Nov 2020 18:19:38 +0100 Subject: [PATCH] BTree: add split_off_range methods --- library/alloc/src/collections/btree/fix.rs | 334 +++++++++++++++++- library/alloc/src/collections/btree/map.rs | 60 ++++ .../alloc/src/collections/btree/map/tests.rs | 28 ++ .../btree/map/tests/remove_range.rs | 297 ++++++++++++++++ .../btree/map/tests/split_off_range.rs | 298 ++++++++++++++++ library/alloc/src/collections/btree/node.rs | 174 +++++++-- library/alloc/src/collections/btree/set.rs | 39 ++ .../alloc/src/collections/btree/set/tests.rs | 12 + library/alloc/src/collections/btree/split.rs | 177 +++++++++- 9 files changed, 1367 insertions(+), 52 deletions(-) create mode 100644 library/alloc/src/collections/btree/map/tests/remove_range.rs create mode 100644 library/alloc/src/collections/btree/map/tests/split_off_range.rs diff --git a/library/alloc/src/collections/btree/fix.rs b/library/alloc/src/collections/btree/fix.rs index c4861817dd05d..d730feca4a2db 100644 --- a/library/alloc/src/collections/btree/fix.rs +++ b/library/alloc/src/collections/btree/fix.rs @@ -1,5 +1,5 @@ use super::map::MIN_LEN; -use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef, Root}; +use super::node::{marker, ForceResult::*, Handle, LeftOrRight, 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 @@ -110,62 +110,374 @@ impl Root { cur_node = last_kv.into_right_child(); } } + + /// Fixes both the left border and the right border. + pub fn fix_both_borders(&mut self) { + self.fix_top(); + while let Internal(root_node) = self.borrow_mut().force() { + if root_node.len() > 1 { + self.fix_both_borders_of_multi_kv_root(); + return; + } + match root_node.try_moving_lone_kv() { + LoneKvResult::MergedIntoEmptyChildren => self.pop_internal_level(), + LoneKvResult::MergedIntoAccompaniedKvIdx(_) => { + self.pop_internal_level(); + self.fix_both_borders_of_multi_kv_root(); + return; + } + LoneKvResult::MovedToAccompanied(_) | LoneKvResult::Required => { + self.fix_both_borders_of_single_kv_root(); + return; + } + } + } + } + + /// Fixes the right border of the left edge and the left border of the right + /// edge, starting from a root node with a single key-value pair, i.e., + /// without siblings and without parent. + pub fn fix_opposite_borders(&mut self) { + assert!(self.len() == 1); + while let Internal(root_node) = self.borrow_mut().force() { + match root_node.try_moving_lone_kv() { + LoneKvResult::MergedIntoEmptyChildren => self.pop_internal_level(), + LoneKvResult::MergedIntoAccompaniedKvIdx(idx) => { + self.pop_internal_level(); + let moved_kv = unsafe { Handle::new_kv(self.borrow_mut(), idx) }; + let fixed = moved_kv.try_fixing_opposite_borders_and_ancestors(); + assert!(fixed); + self.fix_top(); + return; + } + LoneKvResult::MovedToAccompanied(moved_kv) => { + let fixed = moved_kv.try_fixing_opposite_borders_and_ancestors(); + assert!(fixed); + self.fix_top(); + return; + } + LoneKvResult::Required => { + self.fix_opposite_borders_of_single_kv_root(); + return; + } + } + } + } +} + +enum LoneKvResult<'a, K, V> { + MergedIntoEmptyChildren, + MergedIntoAccompaniedKvIdx(usize), + MovedToAccompanied(Handle, K, V, marker::LeafOrInternal>, marker::KV>), + Required, +} + +impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { + /// Detects whether a root with a single key-value pair has underfull + /// children, and tries to honor `MIN_LEN` for both children. + /// If both children are empty, merging one level away is all we can do. + /// If one child is empty and the other filled to capacity, merging is + /// impossible, so we need to steal. + fn try_moving_lone_kv(self) -> LoneKvResult<'a, K, V> { + debug_assert!(self.len() == 1); + let mut here = self.first_kv().consider_for_balancing(); + let (left_len, right_len) = (here.left_child_len(), here.right_child_len()); + if left_len + right_len == 0 { + here.merge(); + LoneKvResult::MergedIntoEmptyChildren + } else if here.can_merge() { + here.merge(); + LoneKvResult::MergedIntoAccompaniedKvIdx(left_len) + } else if left_len < MIN_LEN { + let count = MIN_LEN - left_len; + debug_assert!(right_len >= count + MIN_LEN); + here.bulk_steal_right(count); + let moved_kv = unsafe { Handle::new_kv(here.into_left_child(), left_len) }; + LoneKvResult::MovedToAccompanied(moved_kv) + } else if right_len < MIN_LEN { + let count = MIN_LEN - right_len; + debug_assert!(left_len >= count + MIN_LEN); + here.bulk_steal_left(count); + let moved_kv = unsafe { Handle::new_kv(here.into_right_child(), count - 1) }; + LoneKvResult::MovedToAccompanied(moved_kv) + } else { + LoneKvResult::Required + } + } +} + +impl Root { + /// Fixes both borders of a root with more than one element, which implies + /// the borders never touch. + fn fix_both_borders_of_multi_kv_root(&mut self) { + debug_assert!(self.len() > 1); + self.borrow_mut().first_kv().fix_left_border_of_left_edge(); + self.borrow_mut().last_kv().fix_right_border_of_right_edge(); + self.fix_top(); + } + + /// Fixes both borders of a root with one element and non-empty children. + fn fix_both_borders_of_single_kv_root(&mut self) { + debug_assert!(self.len() == 1); + if let Internal(root_node) = self.borrow_mut().force() { + root_node.first_edge().descend().first_kv().fix_left_border_of_left_edge(); + } else { + unreachable!() + } + if let Internal(root_node) = self.borrow_mut().force() { + root_node.last_edge().descend().last_kv().fix_right_border_of_right_edge(); + } else { + unreachable!() + } + + // Fixing the children may have shrunk them. + if let Internal(root_node) = self.borrow_mut().force() { + root_node.try_moving_lone_kv(); + } else { + unreachable!() + } + self.fix_top(); + } + + /// Fixes opposite borders of a root with one element and non-empty children. + fn fix_opposite_borders_of_single_kv_root(&mut self) { + debug_assert!(self.len() == 1); + if let Internal(root_node) = self.borrow_mut().force() { + root_node.first_edge().descend().last_kv().fix_right_border_of_right_edge(); + } else { + unreachable!() + } + if let Internal(root_node) = self.borrow_mut().force() { + root_node.last_edge().descend().first_kv().fix_left_border_of_left_edge(); + } else { + unreachable!() + } + + // Fixing the children may have shrunk them. + if let Internal(root_node) = self.borrow_mut().force() { + root_node.try_moving_lone_kv(); + } else { + unreachable!() + } + self.fix_top(); + } +} + +impl<'a, K: 'a, V: 'a> Handle, K, V, marker::Internal>, marker::Edge> { + /// Tries to fix the left border of the subtree, using one of the edge's + /// adjacent KVs or the parent. Changes nothing and returns false if + /// the edge has no adjacent KVs and no parent node. + pub fn try_fixing_left_border_and_ancestors(mut self) -> bool { + if unsafe { self.reborrow_mut() }.try_fixing_left_border() { + self.into_node().forget_type().fix_node_and_affected_ancestors(); + true + } else if let Ok(mut parent_edge) = self.into_node().ascend() { + // Since there is no sibling, move on up to the parent, + // an untouched incoming node always able to balance. + let mut edge = match unsafe { parent_edge.reborrow_mut() } + .try_fixing_child_tracking_border(Left(())) + .ok() + .unwrap() + .force() + { + Internal(edge) => edge, + Leaf(_) => unreachable!(), + }; + let fixed = unsafe { edge.reborrow_mut() }.try_fixing_left_border(); + assert!(fixed); + edge.into_node().forget_type().fix_node_through_parent().ok().unwrap(); + parent_edge.into_node().forget_type().fix_node_and_affected_ancestors(); + true + } else { + false + } + } + + fn try_fixing_left_border(self) -> bool { + if let Ok(mut child_edge) = self.try_fixing_child_tracking_border(Left(())) { + loop { + child_edge = match child_edge.force() { + Internal(edge) => edge.try_fixing_child_tracking_border(Left(())).ok().unwrap(), + Leaf(_) => return true, + } + } + } else { + false + } + } +} + +impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInternal>, marker::KV> { + /// Tries to fix the right border of the KV's left edge and the left border + /// of the KV's right edge, using one or both of its sibling KVs, or using + /// its parent. Changes nothing and returns false if the KV has neither + /// siblings nor a parent. + pub fn try_fixing_opposite_borders_and_ancestors(mut self) -> bool { + if unsafe { self.reborrow_mut() }.try_fixing_opposite_borders() { + // `self`'s KV index may be invalid now, but the node is fine. + self.into_node().fix_node_and_affected_ancestors(); + true + } else if let Ok(mut parent_edge) = self.into_node().ascend() { + // Since there is no sibling, move on up to the parent, + // an untouched incoming node always able to balance. + let edge = unsafe { parent_edge.reborrow_mut() } + .try_fixing_child_tracking_border(Left(())) + .ok() + .unwrap(); + let mut kv = edge.right_kv().ok().unwrap(); + let fixed = unsafe { kv.reborrow_mut() }.try_fixing_opposite_borders(); + assert!(fixed); + kv.into_node().fix_node_through_parent().ok().unwrap(); + parent_edge.into_node().forget_type().fix_node_and_affected_ancestors(); + true + } else { + false + } + } + + fn try_fixing_opposite_borders(mut self) -> bool { + let fixed_right_side; + if let Ok(right_kv) = unsafe { self.reborrow_mut() }.right_edge().right_kv() { + right_kv.fix_left_border_of_left_edge(); + fixed_right_side = true + } else { + fixed_right_side = false + } + + // Fix on the left after fixing on the right, because we may move the KV + // that `self` refers to. For the same reason, do not reborrow `self`. + match self.left_edge().left_kv() { + Ok(mut left_kv) => { + unsafe { left_kv.reborrow_mut() }.fix_right_border_of_right_edge(); + if !fixed_right_side { + // Now that the left child is stocked up, we can use it to + // fix the rightmost child in turn. + left_kv.into_node().last_kv().fix_left_border_of_right_edge(); + } + true + } + Err(leftmost_edge) => { + if fixed_right_side { + // Since the right child was stocked up, we can use + // it to fix the leftmost child in turn. + leftmost_edge.into_node().first_kv().fix_right_border_of_left_edge(); + true + } else { + false + } + } + } + } } 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(); + self = internal_kv.fix_left_child(Left(())).into_node().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(); + self = internal_kv.fix_right_child(Right(())).into_node().last_kv(); + debug_assert!(self.reborrow().into_node().len() > MIN_LEN); + } + } + + fn fix_left_border_of_right_edge(mut self) { + while let Internal(internal_kv) = self.force() { + let child_left_edge = internal_kv.fix_right_child(Left(())); + self = match child_left_edge.right_kv() { + Ok(child_right_kv) => { + child_right_kv.fix_left_border_of_left_edge(); + return; + } + Err(child_left_edge) => child_left_edge.into_node().last_kv(), + }; + debug_assert!(self.reborrow().into_node().len() > MIN_LEN); + } + } + + fn fix_right_border_of_left_edge(mut self) { + while let Internal(internal_kv) = self.force() { + let child_right_edge = internal_kv.fix_left_child(Right(())); + self = match child_right_edge.left_kv() { + Ok(child_left_kv) => { + child_left_kv.fix_right_border_of_right_edge(); + return; + } + Err(child_right_edge) => child_right_edge.into_node().first_kv(), + }; debug_assert!(self.reborrow().into_node().len() > MIN_LEN); } } } +impl<'a, K: 'a, V: 'a> Handle, K, V, marker::Internal>, marker::Edge> { + /// Fixes the child only, using one of the adjacent parent KVs. + /// Returns the edge where the left child of the original child ended up. + fn try_fixing_child_tracking_border( + self, + track_border: LeftOrRight<()>, + ) -> Result, K, V, marker::LeafOrInternal>, marker::Edge>, Self> + { + match self.left_kv() { + Ok(left_kv) => Ok(left_kv.fix_right_child(track_border)), + Err(edge) => match edge.right_kv() { + Ok(right_kv) => Ok(right_kv.fix_left_child(track_border)), + Err(edge) => Err(edge), + }, + } + } +} + 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> { + /// Returns the edge where the tracked border of the original left child ended up. + fn fix_left_child( + self, + track_border: LeftOrRight<()>, + ) -> Handle, K, V, marker::LeafOrInternal>, marker::Edge> { 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); + let track_edge_idx = if matches!(track_border, Left(_)) { 0 } else { left_len }; if internal_kv.can_merge() { - internal_kv.merge_tracking_child() + internal_kv.merge_tracking_child_edge(Left(track_edge_idx)) } 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() + unsafe { Handle::new_edge(internal_kv.into_left_child(), track_edge_idx) } } } /// 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> { + /// Returns the edge where the tracked border of the original right child ended up. + fn fix_right_child( + self, + track_border: LeftOrRight<()>, + ) -> Handle, K, V, marker::LeafOrInternal>, marker::Edge> { 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); + let track_edge_idx = if matches!(track_border, Left(_)) { 0 } else { right_len }; if internal_kv.can_merge() { - internal_kv.merge_tracking_child() + internal_kv.merge_tracking_child_edge(Right(track_edge_idx)) } 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() + unsafe { Handle::new_edge(internal_kv.into_right_child(), count + track_edge_idx) } } } } diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index cdb961d4cfbc5..5fc335755aa95 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1204,6 +1204,66 @@ impl BTreeMap { BTreeMap { root: Some(right_root), length: right_len } } + /// Splits the collection into two. Returns a new collection with all keys in the given range. + /// + /// # Panics + /// + /// Panics if range `start > end`. + /// Panics if range `start == end` and both bounds are `Excluded`. + /// May panic if the [`Ord`] implementation of type `T` is ill-defined, + /// either because it does not form a total order or because it does not + /// correspond to the [`Ord`] implementation of type `K`. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(btree_drain)] + /// use std::collections::BTreeMap; + /// + /// let mut a = BTreeMap::new(); + /// a.insert(1, "a"); + /// a.insert(2, "b"); + /// a.insert(3, "c"); + /// a.insert(17, "d"); + /// a.insert(41, "e"); + /// + /// let b = a.split_off_range(&3..&33); + /// + /// assert_eq!(a.len(), 3); + /// assert_eq!(b.len(), 2); + /// + /// assert_eq!(a[&1], "a"); + /// assert_eq!(a[&2], "b"); + /// assert_eq!(a[&41], "e"); + /// + /// assert_eq!(b[&3], "c"); + /// assert_eq!(b[&17], "d"); + /// ``` + #[unstable(feature = "btree_drain", issue = "81074")] + pub fn split_off_range(&mut self, range: R) -> Self + where + T: Ord, + K: Borrow + Ord, + R: RangeBounds, + { + if self.is_empty() { + return Self::new(); + } + + let total_num = self.length; + let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty + + let mut right_root = left_root.split_off_range(range); + right_root.fix_both_borders(); + + let (new_left_len, right_len) = Root::calc_split_length(total_num, &left_root, &right_root); + self.length = new_left_len; + + BTreeMap { root: Some(right_root), length: right_len } + } + /// Creates an iterator that visits all elements (key-value pairs) in /// ascending key order and uses a closure to determine if an element should /// be removed. If the closure returns `true`, the element is removed from diff --git a/library/alloc/src/collections/btree/map/tests.rs b/library/alloc/src/collections/btree/map/tests.rs index b39b5409ae44f..3f0519358a35e 100644 --- a/library/alloc/src/collections/btree/map/tests.rs +++ b/library/alloc/src/collections/btree/map/tests.rs @@ -17,6 +17,8 @@ use std::ops::RangeBounds; use std::panic::{catch_unwind, AssertUnwindSafe}; use std::sync::atomic::{AtomicUsize, Ordering::SeqCst}; +mod split_off_range; + // Capacity of a tree with a single level, // i.e., a tree who's root is a leaf node at height 0. const NODE_CAPACITY: usize = node::CAPACITY; @@ -31,6 +33,12 @@ const MIN_INSERTS_HEIGHT_1: usize = NODE_CAPACITY + 1; // It's not the minimum size: removing an element from such a tree does not always reduce height. const MIN_INSERTS_HEIGHT_2: usize = 89; +// Like MIN_INSERTS_HEIGHT_2, with an additional internal level. +const MIN_INSERTS_HEIGHT_3: usize = 628; + +// Like MIN_INSERTS_HEIGHT_3, with an additional internal level. +const MIN_INSERTS_HEIGHT_4: usize = 4401; + // Gathers all references from a mutable iterator and makes sure Miri notices if // using them is dangerous. fn test_all_refs<'a, T: 'a>(dummy: &mut T, iter: impl Iterator) { @@ -173,6 +181,26 @@ fn test_levels() { // - 5 elements in right child's last grandchild assert_eq!(map.height(), Some(2)); assert_eq!(map.len(), MIN_INSERTS_HEIGHT_2, "{}", map.dump_keys()); + + if cfg!(miri) { + // Miri is too slow + return; + } + while map.height() == Some(2) { + let last_key = *map.last_key_value().unwrap().0; + map.insert(last_key + 1, ()); + } + map.check(); + assert_eq!(map.height(), Some(3)); + assert_eq!(map.len(), MIN_INSERTS_HEIGHT_3); + + while map.height() == Some(3) { + let last_key = *map.last_key_value().unwrap().0; + map.insert(last_key + 1, ()); + } + map.check(); + assert_eq!(map.height(), Some(4)); + assert_eq!(map.len(), MIN_INSERTS_HEIGHT_4); } // Ensures the testing infrastructure usually notices order violations. diff --git a/library/alloc/src/collections/btree/map/tests/remove_range.rs b/library/alloc/src/collections/btree/map/tests/remove_range.rs new file mode 100644 index 0000000000000..bf88b5927bc48 --- /dev/null +++ b/library/alloc/src/collections/btree/map/tests/remove_range.rs @@ -0,0 +1,297 @@ +use super::*; +use std::cmp::{max, min}; + +#[test] +fn empty() { + let mut map: BTreeMap = BTreeMap::new(); + map.remove_range(..); + assert!(map.is_empty()); + map.check(); +} + +// Drop the iterator, where most test cases consume it entirely. +#[test] +fn dropped_keeping_all() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.remove_range(..0); + assert!(map.keys().copied().eq(0..3)); + map.check(); +} + +// Drop the iterator, where most test cases consume it entirely. +#[test] +fn dropped_removing_all() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + map.remove_range(..); + assert!(map.is_empty()); + map.check(); +} + +#[test] +fn consumed_keeping_all() { + let pairs = (0..3).map(|i| (i, ())); + let mut map: BTreeMap<_, _> = pairs.collect(); + assert_eq!(map.remove_range(..0), 0); + assert!(map.keys().copied().eq(0..3)); + map.check(); +} + +fn test_size_range>( + size: usize, + height: usize, + compact: bool, + range: R, + keep: usize, +) { + let mut map: BTreeMap<_, _> = if compact { + (0..size).map(|i| (i, ())).collect() + } else { + let mut map = BTreeMap::new(); + for i in 0..size { + map.insert(i, ()); + } + map + }; + assert_eq!(map.len(), size); + assert_eq!(map.height(), Some(height), "{}", map.dump_keys()); + let removed = map.remove_range(range); + assert_eq!(removed, size - keep); + assert_eq!(map.len(), keep); + map.check(); +} + +fn test_size_keeping_n(size: usize, height: usize, compact: bool, keep: usize) { + for doomed_start in 0..keep + 1 { + test_size_range(size, height, compact, doomed_start..(doomed_start + size - keep), keep); + } +} + +fn test_size_all(size: usize, height: usize, compact: bool) { + for keep in 0..size + 1 { + test_size_keeping_n(size, height, compact, keep) + } +} + +fn test_size_some(size: usize, height: usize, compact: bool) { + test_size_keeping_n(size, height, compact, 0); + test_size_keeping_n(size, height, compact, 1); + test_size_keeping_n(size, height, compact, 2); + test_size_keeping_n(size, height, compact, size / 4); + test_size_keeping_n(size, height, compact, size / 2); + test_size_keeping_n(size, height, compact, size - 2); + test_size_keeping_n(size, height, compact, size - 1); +} + +#[test] +fn height_0_underfull_all() { + test_size_all(3, 0, false) +} + +#[test] +fn height_0_max_some() { + test_size_some(NODE_CAPACITY, 0, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_0_max_all() { + test_size_all(NODE_CAPACITY, 0, false) +} + +#[test] +fn height_1_min_keeping_0() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 0) +} + +#[test] +fn height_1_min_keeping_1() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 1) +} + +#[test] +fn height_1_min_keeping_2() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 2) +} + +#[test] +fn height_1_min_keeping_7() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 7) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_1_min_all() { + test_size_all(MIN_INSERTS_HEIGHT_1, 1, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_1_more_all() { + for size in MIN_INSERTS_HEIGHT_1 + 1..MIN_INSERTS_HEIGHT_2 { + test_size_all(size, 1, false) + } +} + +#[test] +fn height_2_min_keeping_0() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_2, 2, false, 0) +} + +#[test] +fn height_2_min_keeping_1() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_2, 2, false, 1) +} + +#[test] +fn height_2_min_keeping_12_left() { + test_size_range(MIN_INSERTS_HEIGHT_2, 2, false, 0..77, 12); +} + +#[test] +fn height_2_min_keeping_12_mid() { + test_size_range(MIN_INSERTS_HEIGHT_2, 2, false, 6..83, 12); +} + +#[test] +fn height_2_min_keeping_12_right() { + test_size_range(MIN_INSERTS_HEIGHT_2, 2, false, 12..89, 12) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_2_min_all() { + test_size_all(MIN_INSERTS_HEIGHT_2, 2, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn height_2_more_some() { + for size in MIN_INSERTS_HEIGHT_2 + 1..MIN_INSERTS_HEIGHT_3 { + println!("size {}", size); + test_size_some(size, 2, false) + } +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_3_min_some() { + test_size_some(MIN_INSERTS_HEIGHT_3, 3, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn height_3_min_all() { + test_size_all(MIN_INSERTS_HEIGHT_3, 3, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn height_4_min_some() { + test_size_some(MIN_INSERTS_HEIGHT_4, 4, false) +} + +#[test] +fn size_143_compact_keeping_1() { + test_size_keeping_n(143, 1, true, 1) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_143_compact_all() { + test_size_all(143, 1, true) +} + +#[test] +fn size_144_compact_keeping_1() { + test_size_keeping_n(144, 2, true, 1) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_144_compact_all() { + test_size_all(144, 2, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_1727_compact_some() { + test_size_some(1727, 2, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_1728_compact_some() { + test_size_some(1728, 3, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn size_20735_compact_some() { + test_size_some(20735, 3, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn size_20736_compact_some() { + test_size_some(20736, 4, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn sub_size_143_compact_some() { + for size in NODE_CAPACITY + 1..143 { + test_size_some(size, 1, true) + } +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn sub_size_1727_compact_some() { + for size in (144 + 1..1727).step_by(10) { + test_size_some(size, 2, true) + } +} + +#[test] +fn random_1() { + let mut rng = DeterministicRng::new(); + for _ in 0..if cfg!(miri) { 1 } else { 140 } { + let size = rng.next() as usize % 1024; + let mut map: BTreeMap<_, ()> = BTreeMap::new(); + for _ in 0..size { + map.insert(rng.next(), ()); + } + assert_eq!(map.len(), size); + let (x, y) = (rng.next(), rng.next()); + let bounds = min(x, y)..max(x, y); + let removed = map.remove_range(bounds.clone()); + assert_eq!(removed + map.len(), size); + map.check(); + assert!(!map.into_keys().any(|k| bounds.contains(&k))); + } +} + +#[test] +fn drop_panic_leak() { + let a = CrashTestDummy::new(0); + let b = CrashTestDummy::new(1); + let c = CrashTestDummy::new(2); + let mut map = BTreeMap::new(); + map.insert(a.spawn(Panic::Never), ()); + map.insert(b.spawn(Panic::InDrop), ()); + map.insert(c.spawn(Panic::Never), ()); + + catch_unwind(move || drop(map.remove_range(..))).unwrap_err(); + + assert_eq!(a.dropped(), 1); + assert_eq!(b.dropped(), 1); + assert_eq!(c.dropped(), 1); +} diff --git a/library/alloc/src/collections/btree/map/tests/split_off_range.rs b/library/alloc/src/collections/btree/map/tests/split_off_range.rs new file mode 100644 index 0000000000000..0bca5a873c7a2 --- /dev/null +++ b/library/alloc/src/collections/btree/map/tests/split_off_range.rs @@ -0,0 +1,298 @@ +use super::*; +use std::cmp::{max, min}; + +#[test] +fn empty() { + let mut map: BTreeMap = BTreeMap::new(); + map.split_off_range(..); + assert!(map.is_empty()); + map.check(); +} + +// Drop the iterator, where most test cases consume it entirely. +#[test] +fn dropped_keeping_all() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.collect(); + map.split_off_range(..0); + assert!(map.keys().copied().eq(0..3)); + map.check(); +} + +// Drop the iterator, where most test cases consume it entirely. +#[test] +fn dropped_removing_all() { + let pairs = (0..3).map(|i| (i, i)); + let mut map: BTreeMap<_, _> = pairs.clone().collect(); + map.split_off_range(..); + assert!(map.is_empty()); + map.check(); +} + +#[test] +fn consumed_keeping_all() { + let pairs = (0..3).map(|i| (i, ())); + let mut map: BTreeMap<_, _> = pairs.collect(); + assert!(map.split_off_range(..0).eq(&BTreeMap::new())); + assert!(map.keys().copied().eq(0..3)); + map.check(); +} + +fn test_size_range>( + size: usize, + height: usize, + compact: bool, + range: R, + keep: usize, +) { + let mut map: BTreeMap<_, _> = if compact { + (0..size).map(|i| (i, ())).collect() + } else { + let mut map = BTreeMap::new(); + for i in 0..size { + map.insert(i, ()); + } + map + }; + assert_eq!(map.len(), size); + assert_eq!(map.height(), Some(height), "{}", map.dump_keys()); + let split_off = map.split_off_range(range); + assert_eq!(split_off.len(), size - keep); + assert_eq!(map.len(), keep); + split_off.check(); + map.check(); +} + +// Example of a way to debug these test cases. +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn dbg() { + let size = 181; + let range = 7..98; + let mut map: BTreeMap<_, _> = (0..size).map(|i| (i, ())).collect(); + let mut root = map.root.take().unwrap(); + println!("in: {}\n", root.reborrow().dump_keys()); + let drained = root.split_off_range(range); + println!("kept: {}\n", root.reborrow().dump_keys()); + println!("drained: {}\n", drained.reborrow().dump_keys()); +} + +fn test_size_keeping_n(size: usize, height: usize, compact: bool, keep: usize) { + for doomed_start in 0..keep + 1 { + test_size_range(size, height, compact, doomed_start..(doomed_start + size - keep), keep); + } +} + +fn test_size_all(size: usize, height: usize, compact: bool) { + for keep in 0..size + 1 { + test_size_keeping_n(size, height, compact, keep) + } +} + +fn test_size_some(size: usize, height: usize, compact: bool) { + test_size_keeping_n(size, height, compact, 0); + test_size_keeping_n(size, height, compact, 1); + test_size_keeping_n(size, height, compact, 2); + test_size_keeping_n(size, height, compact, size / 4); + test_size_keeping_n(size, height, compact, size / 2); + test_size_keeping_n(size, height, compact, size - 2); + test_size_keeping_n(size, height, compact, size - 1); +} + +#[test] +fn height_0_underfull_all() { + test_size_all(3, 0, false) +} + +#[test] +fn height_0_max_some() { + test_size_some(NODE_CAPACITY, 0, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_0_max_all() { + test_size_all(NODE_CAPACITY, 0, false) +} + +#[test] +fn height_1_min_keeping_0() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 0) +} + +#[test] +fn height_1_min_keeping_1() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 1) +} + +#[test] +fn height_1_min_keeping_2() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 2) +} + +#[test] +fn height_1_min_keeping_7() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_1, 1, false, 7) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_1_min_all() { + test_size_all(MIN_INSERTS_HEIGHT_1, 1, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_1_more_all() { + for size in MIN_INSERTS_HEIGHT_1 + 1..MIN_INSERTS_HEIGHT_2 { + test_size_all(size, 1, false) + } +} + +#[test] +fn height_2_min_keeping_0() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_2, 2, false, 0) +} + +#[test] +fn height_2_min_keeping_1() { + test_size_keeping_n(MIN_INSERTS_HEIGHT_2, 2, false, 1) +} + +#[test] +fn height_2_min_keeping_12_left() { + test_size_range(MIN_INSERTS_HEIGHT_2, 2, false, 0..77, 12); +} + +#[test] +fn height_2_min_keeping_12_mid() { + test_size_range(MIN_INSERTS_HEIGHT_2, 2, false, 6..83, 12); +} + +#[test] +fn height_2_min_keeping_12_right() { + test_size_range(MIN_INSERTS_HEIGHT_2, 2, false, 12..89, 12) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_2_min_all() { + test_size_all(MIN_INSERTS_HEIGHT_2, 2, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn height_2_more_some() { + for size in MIN_INSERTS_HEIGHT_2 + 1..MIN_INSERTS_HEIGHT_3 { + println!("size {}", size); + test_size_some(size, 2, false) + } +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn height_3_min_some() { + test_size_some(MIN_INSERTS_HEIGHT_3, 3, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn height_3_min_all() { + test_size_all(MIN_INSERTS_HEIGHT_3, 3, false) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn height_4_min_some() { + test_size_some(MIN_INSERTS_HEIGHT_4, 4, false) +} + +#[test] +fn size_143_compact_keeping_1() { + test_size_keeping_n(143, 1, true, 1) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_143_compact_all() { + test_size_all(143, 1, true) +} + +#[test] +fn size_144_compact_keeping_1() { + test_size_keeping_n(144, 2, true, 1) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_144_compact_all() { + test_size_all(144, 2, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_1727_compact_some() { + test_size_some(1727, 2, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn size_1728_compact_some() { + test_size_some(1728, 3, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn size_20735_compact_some() { + test_size_some(20735, 3, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn size_20736_compact_some() { + test_size_some(20736, 4, true) +} + +#[cfg(not(miri))] // Miri is too slow +#[test] +fn sub_size_143_compact_some() { + for size in NODE_CAPACITY + 1..143 { + test_size_some(size, 1, true) + } +} + +#[cfg(not(miri))] // Miri is too slow +#[ignore] +#[test] +fn sub_size_1727_compact_some() { + for size in (144 + 1..1727).step_by(10) { + test_size_some(size, 2, true) + } +} + +#[test] +fn random_1() { + let mut rng = DeterministicRng::new(); + for _ in 0..if cfg!(miri) { 1 } else { 140 } { + let size = rng.next() as usize % 1024; + let mut map: BTreeMap<_, ()> = BTreeMap::new(); + for _ in 0..size { + map.insert(rng.next(), ()); + } + assert_eq!(map.len(), size); + let (x, y) = (rng.next(), rng.next()); + let bounds = min(x, y)..max(x, y); + let split_off = map.split_off_range(bounds.clone()); + assert_eq!(split_off.len() + map.len(), size); + split_off.check(); + map.check(); + assert!(split_off.into_keys().all(|k| bounds.contains(&k))); + assert!(!map.into_keys().any(|k| bounds.contains(&k))); + } +} diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index dfce98f97bd44..c4fd62f05fb2a 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -213,7 +213,7 @@ unsafe impl Send for NodeRef unsafe impl Send for NodeRef {} impl NodeRef { - fn new_leaf() -> Self { + pub fn new_leaf() -> Self { Self::from_new_leaf(LeafNode::new()) } @@ -223,7 +223,7 @@ impl NodeRef { } impl NodeRef { - fn new_internal(child: Root) -> Self { + pub fn new_internal(child: Root) -> Self { let mut new_node = unsafe { InternalNode::new() }; new_node.edges[0].write(child.node); unsafe { NodeRef::from_new_internal(new_node, child.height + 1) } @@ -418,7 +418,7 @@ impl<'a, K, V, Type> NodeRef, K, V, Type> { // FIXME(@gereeter) consider adding yet another type parameter to `NodeRef` // that restricts the use of navigation methods on reborrowed pointers, // preventing this unsafety. - unsafe fn reborrow_mut(&mut self) -> NodeRef, K, V, Type> { + pub unsafe fn reborrow_mut(&mut self) -> NodeRef, K, V, Type> { NodeRef { height: self.height, node: self.node, _marker: PhantomData } } @@ -651,6 +651,34 @@ impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { } } +impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { + /// Removes a key-value pair from the end of the node and returns the pair. + /// Also removes the edge that was to the right of that pair and, if the node + /// is internal, returns the orphaned subtree that this edge owned. + pub fn pop(&mut self) -> (K, V, Option>) { + let idx = self.len().checked_sub(1).expect("attempt to pop from empty node"); + + unsafe { + let key = self.key_area_mut(idx).assume_init_read(); + let val = self.val_area_mut(idx).assume_init_read(); + let edge = match self.reborrow_mut().force() { + ForceResult::Leaf(_) => None, + ForceResult::Internal(mut internal) => { + let node = internal.edge_area_mut(idx + 1).assume_init_read(); + let mut edge = Root { node, height: internal.height - 1, _marker: PhantomData }; + // Currently, clearing the parent link is superfluous, because we will + // insert the node elsewhere and set its parent link again. + edge.clear_parent_link(); + Some(edge) + } + }; + + *self.len_mut() -= 1; + (key, val, edge) + } + } +} + impl NodeRef { /// Removes any static information asserting that this node is a `Leaf` node. pub fn forget_type(self) -> NodeRef { @@ -1315,23 +1343,28 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> { } /// Merges the parent's key-value pair and both adjacent child nodes into - /// the left child node and returns the shrunk parent node. + /// the left child node. + /// + /// Panics unless we `.can_merge()`. + pub fn merge(self) { + self.do_merge(|_, _| ()) + } + + /// Performs `merge()` and returns the shrunk parent node. /// /// Panics unless we `.can_merge()`. pub fn merge_tracking_parent(self) -> NodeRef, K, V, marker::Internal> { self.do_merge(|parent, _child| parent) } - /// Merges the parent's key-value pair and both adjacent child nodes into - /// the left child node and returns that child node. + /// Performs `merge()` and returns the expanded child node. /// /// Panics unless we `.can_merge()`. pub fn merge_tracking_child(self) -> NodeRef, K, V, marker::LeafOrInternal> { self.do_merge(|_parent, child| child) } - /// Merges the parent's key-value pair and both adjacent child nodes into - /// the left child node and returns the edge handle in that child node + /// Performs `merge()` and returns the edge handle in the child node /// where the tracked child edge ended up, /// /// Panics unless we `.can_merge()`. @@ -1573,45 +1606,108 @@ impl<'a, K, V> Handle, K, V, marker::LeafOrInternal>, ma &mut self, right: &mut NodeRef, K, V, marker::LeafOrInternal>, ) { - unsafe { - let new_left_len = self.idx; - let mut left_node = self.reborrow_mut().into_node(); - let old_left_len = left_node.len(); - - let new_right_len = old_left_len - new_left_len; - let mut right_node = right.reborrow_mut(); - - assert!(right_node.len() == 0); - assert!(left_node.height == right_node.height); + self.move_region(None, right) + } - if new_right_len > 0 { - *left_node.len_mut() = new_left_len as u16; - *right_node.len_mut() = new_right_len as u16; + /// Moves a region from `self` to the back of another node. The region + /// consists of the key-value pairs and edges right of `self`, up to and + /// including the last edge specified by index. + /// Any existing key-value pairs and edges in `dst_node` remain unchanged. + fn move_region( + &mut self, + last_src_edge_idx: Option, + dst_node: &mut NodeRef, K, V, marker::LeafOrInternal>, + ) { + let mut src_node = unsafe { self.node.reborrow_mut() }; + let mut dst_node = unsafe { dst_node.reborrow_mut() }; + assert!(dst_node.height == src_node.height); + let old_src_len = src_node.len(); + let old_dst_len = dst_node.len(); + let first_src_edge_idx = self.idx; + let last_src_edge_idx = last_src_edge_idx.unwrap_or(old_src_len); + assert!(last_src_edge_idx <= old_src_len); + assert!(last_src_edge_idx >= first_src_edge_idx); + let infix_len = last_src_edge_idx - first_src_edge_idx; + let new_src_len = old_src_len - infix_len; + let new_dst_len = old_dst_len + infix_len; + assert!(new_dst_len <= CAPACITY); + + *src_node.len_mut() = new_src_len as u16; + *dst_node.len_mut() = new_dst_len as u16; - move_to_slice( - left_node.key_area_mut(new_left_len..old_left_len), - right_node.key_area_mut(..new_right_len), - ); - move_to_slice( - left_node.val_area_mut(new_left_len..old_left_len), - right_node.val_area_mut(..new_right_len), - ); - match (left_node.force(), right_node.force()) { - (ForceResult::Internal(mut left), ForceResult::Internal(mut right)) => { - move_to_slice( - left.edge_area_mut(new_left_len + 1..old_left_len + 1), - right.edge_area_mut(1..new_right_len + 1), - ); - right.correct_childrens_parent_links(1..new_right_len + 1); - } - (ForceResult::Leaf(_), ForceResult::Leaf(_)) => {} - _ => unreachable!(), + unsafe { + move_to_slice( + src_node.key_area_mut(first_src_edge_idx..last_src_edge_idx), + dst_node.key_area_mut(old_dst_len..new_dst_len), + ); + move_to_slice( + src_node.val_area_mut(first_src_edge_idx..last_src_edge_idx), + dst_node.val_area_mut(old_dst_len..new_dst_len), + ); + slice_shl(src_node.key_area_mut(first_src_edge_idx..old_src_len), infix_len); + slice_shl(src_node.val_area_mut(first_src_edge_idx..old_src_len), infix_len); + match (src_node.force(), dst_node.force()) { + (ForceResult::Internal(mut src_node), ForceResult::Internal(mut dst_node)) => { + move_to_slice( + src_node.edge_area_mut(first_src_edge_idx + 1..last_src_edge_idx + 1), + dst_node.edge_area_mut(old_dst_len + 1..new_dst_len + 1), + ); + slice_shl( + src_node.edge_area_mut(first_src_edge_idx + 1..old_src_len + 1), + infix_len, + ); + src_node + .correct_childrens_parent_links(first_src_edge_idx + 1..new_src_len + 1); + dst_node.correct_childrens_parent_links(old_dst_len + 1..new_dst_len + 1); } + (ForceResult::Leaf(_), ForceResult::Leaf(_)) => {} + _ => unreachable!(), } } } } +impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::Edge> { + pub fn move_infix( + &mut self, + last_src_edge_idx: usize, + dst_node: &mut NodeRef, K, V, marker::Leaf>, + ) { + let mut src_node = unsafe { self.reborrow_mut().forget_node_type() }; + let mut dst_node = unsafe { dst_node.reborrow_mut().forget_type() }; + src_node.move_region(Some(last_src_edge_idx), &mut dst_node) + } +} + +impl<'a, K, V> Handle, K, V, marker::Internal>, marker::Edge> { + pub fn move_infix( + &mut self, + last_src_edge_idx: usize, + dst_node: &mut NodeRef, K, V, marker::Internal>, + ) { + let mut src_node = unsafe { self.reborrow_mut().forget_node_type() }; + let mut dst_node = unsafe { dst_node.reborrow_mut().forget_type() }; + src_node.move_region(Some(last_src_edge_idx), &mut dst_node) + } +} + +impl<'a, K: 'a, V: 'a> Handle, K, V, marker::Internal>, marker::Edge> { + /// Replaces a single edge in `self` and returns the orphaned previous edge. + pub fn replace_edge(&mut self, replacement_edge: Root) -> Root { + assert!(replacement_edge.height == self.node.height - 1); + unsafe { + let node = mem::replace( + self.node.edge_area_mut(self.idx).assume_init_mut(), + replacement_edge.node, + ); + self.reborrow_mut().correct_parent_link(); + let mut old_edge = Root { node, height: self.node.height - 1, _marker: PhantomData }; + old_edge.clear_parent_link(); + old_edge + } + } +} + pub enum ForceResult { Leaf(Leaf), Internal(Internal), diff --git a/library/alloc/src/collections/btree/set.rs b/library/alloc/src/collections/btree/set.rs index 31df4e98ed746..37829493fafc5 100644 --- a/library/alloc/src/collections/btree/set.rs +++ b/library/alloc/src/collections/btree/set.rs @@ -968,6 +968,45 @@ impl BTreeSet { BTreeSet { map: self.map.split_off(value) } } + /// Splits the collection into two. Returns a new collection with all elements in the given range. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(btree_drain)] + /// use std::collections::BTreeSet; + /// + /// let mut a = BTreeSet::new(); + /// a.insert(1); + /// a.insert(2); + /// a.insert(3); + /// a.insert(17); + /// a.insert(41); + /// + /// let b = a.split_off_range(&3..&33); + /// + /// assert_eq!(a.len(), 3); + /// assert_eq!(b.len(), 2); + /// + /// assert!(a.contains(&1)); + /// assert!(a.contains(&2)); + /// assert!(a.contains(&41)); + /// + /// assert!(b.contains(&3)); + /// assert!(b.contains(&17)); + /// ``` + #[unstable(feature = "btree_drain", issue = "81074")] + pub fn split_off_range(&mut self, range: R) -> Self + where + K: Ord, + T: Borrow + Ord, + R: RangeBounds, + { + BTreeSet { map: self.map.split_off_range(range) } + } + /// Creates an iterator that visits all elements in ascending order and /// uses a closure to determine if an element should be removed. /// diff --git a/library/alloc/src/collections/btree/set/tests.rs b/library/alloc/src/collections/btree/set/tests.rs index 7390ff5a59c8c..451a8277b3518 100644 --- a/library/alloc/src/collections/btree/set/tests.rs +++ b/library/alloc/src/collections/btree/set/tests.rs @@ -770,6 +770,18 @@ fn test_split_off_large_random_sorted() { assert!(right.into_iter().eq(data.into_iter().filter(|x| *x >= key))); } +#[test] +fn test_split_off_range() { + let mut set: BTreeSet<_> = (1..=3).collect(); + let drained = set.split_off_range(2..=2); + assert_eq!(drained.len(), 1); + assert_eq!(drained.first(), Some(&2)); + assert_eq!(drained.last(), Some(&2)); + assert_eq!(set.len(), 2); + assert_eq!(set.first(), Some(&1)); + assert_eq!(set.last(), Some(&3)); +} + #[test] fn from_array() { let set = BTreeSet::from([1, 2, 3, 4]); diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs index bec495a72a4a8..5352ae5c3bbef 100644 --- a/library/alloc/src/collections/btree/split.rs +++ b/library/alloc/src/collections/btree/split.rs @@ -1,6 +1,9 @@ -use super::node::{ForceResult::*, Root}; -use super::search::SearchResult::*; +use super::node::LeftOrRight::{self, *}; +use super::node::{marker, ForceResult::*, Handle, NodeRef, Root}; +use super::search::{SearchBound, SearchResult::*}; use core::borrow::Borrow; +use core::ops::RangeBounds; +use core::ptr; impl Root { /// Calculates the length of both trees that result from splitting up @@ -61,6 +64,108 @@ impl Root { right_root } + /// Splits off a tree with the key-value pairs contained within a range. + /// The returned tree respects the tree invariants of the `node` module, + /// and has keys in ascending order, but does not care about the other + /// tree invariants of `BTreeMap`. Invoke its `fix_both_borders` to make + /// it a valid root for a `BTreeMap`. + pub fn split_off_range(&mut self, range: R) -> Self + where + Q: ?Sized + Ord, + K: Borrow, + R: RangeBounds, + { + let (node, lower_edge_idx, upper_edge_idx, lower_child_bound, upper_child_bound) = + match self.borrow_mut().search_tree_for_bifurcation(&range) { + Ok(found) => found, + Err(_) => return Root::new(), + }; + debug_assert!(lower_edge_idx < upper_edge_idx); + + // We're at the highest node that contains a non-empty sequence of + // key-value pairs within the range. Move out those pairs, and any edge + // in between them, to form the root of the new tree. + // Then descend further along the left and right border, on each level + // splitting up one node along the lower bound (start of the range), and + // splitting up one node along the upper bound (end of the range), + // resulting in: + // - Lower stems, the parts remaining along the lower bound. + // - Lower clippings, the parts split off along the lower bound. + // - Upper stems, the parts remaining along the upper bound, that belong + // in the new tree. + // - Upper clippings, the parts split off along the upper bound, that + // belong in `self`. + // Each of those can easily be underfull or empty, which makes it hard + // to preserve the `BTreeMap` invariants. + // + // In addition to that, repairing `self` is hard because we're left with + // with two edges and no key-value pair to glue them together. Therefore, + // we dig up a key-value pair from one of the edges. + // + // Constructing the new tree is merely gluing two outer edges to the + // moved sequence of key-value pairs. + match node.force() { + Leaf(leaf) => { + let mut leaf_edge = unsafe { Handle::new_edge(leaf, lower_edge_idx) }; + let mut new_root = NodeRef::new_leaf(); + leaf_edge.move_infix(upper_edge_idx, &mut new_root.borrow_mut()); + leaf_edge.into_node().forget_type().fix_node_and_affected_ancestors(); + self.fix_top(); + new_root.forget_type() + } + Internal(node) => { + let child_height = node.height() - 1; + let mut lower_edge = unsafe { Handle::new_edge(ptr::read(&node), lower_edge_idx) }; + let upper_edge = unsafe { Handle::new_edge(node, upper_edge_idx) }; + let lower_child = unsafe { ptr::read(&lower_edge) }.descend(); + let upper_child = unsafe { ptr::read(&upper_edge) }.descend(); + let mut lower_clippings = Root::new_pillar(child_height); + let mut upper_clippings = Root::new_pillar(child_height); + let middle_kv = lower_child.split_off_bound( + lower_clippings.borrow_mut(), + lower_child_bound, + Left(()), + ); + upper_child.split_off_bound( + upper_clippings.borrow_mut(), + upper_child_bound, + Right(()), + ); + // We keep the edge left of the infix (minus the part split off into + // `upper_clippings`), move out the infix and the edge right of + // the infix (minus the part split off), and replace it with + // `upper_clippings`. But we need a KV to bridge the edges. + let mut dst_root = NodeRef::new_internal(lower_clippings); + if let Some(middle_kv) = middle_kv { + let mut lower_kv = lower_edge.right_kv().ok().unwrap(); + let first_kv = lower_kv.replace_kv(middle_kv.0, middle_kv.1); + let mut lower_edge_plus_1 = unsafe { lower_kv.reborrow_mut() }.right_edge(); + let upper_stems = lower_edge_plus_1.replace_edge(upper_clippings); + dst_root.borrow_mut().push(first_kv.0, first_kv.1, upper_stems); + lower_edge_plus_1.move_infix(upper_edge.idx(), &mut dst_root.borrow_mut()); + if lower_kv.forget_node_type().try_fixing_opposite_borders_and_ancestors() { + self.fix_top(); + } else { + self.fix_opposite_borders(); + } + } else { + // There's no KV in the remainder of the left border, so instead we + // discard that empty remainder. + let lower_stems = + unsafe { lower_edge.reborrow_mut() }.replace_edge(upper_clippings); + lower_stems.deallocate_pillar(); + lower_edge.move_infix(upper_edge.idx(), &mut dst_root.borrow_mut()); + if lower_edge.try_fixing_left_border_and_ancestors() { + self.fix_top(); + } else { + self.fix_left_border(); + } + } + dst_root.forget_type() + } + } + } + /// Creates a tree consisting of empty nodes. fn new_pillar(height: usize) -> Self { let mut root = Root::new(); @@ -69,4 +174,72 @@ impl Root { } root } + + /// Destroys a tree consisting of empty nodes. + fn deallocate_pillar(mut self) { + debug_assert!(self.reborrow().calc_length() == 0); + while self.height() > 0 { + self.pop_internal_level(); + } + unsafe { self.into_dying().deallocate_and_ascend() }; + } +} + +impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::LeafOrInternal> { + /// Splits off a subtree at a given left (lower) or right (upper) bound. + /// For the left bound, also splits off the greatest key-value pair below + /// the bound, if any, and returns it. + fn split_off_bound<'q, Q>( + mut self, + mut dst_node: Self, + mut bound: SearchBound<&'q Q>, + border: LeftOrRight<()>, + ) -> Option<(K, V)> + where + Q: ?Sized + Ord, + K: Borrow, + { + let top_height = self.height(); + loop { + let (mut next_edge, next_bound) = match border { + Left(_) => self.find_lower_bound_edge(bound), + Right(_) => self.find_upper_bound_edge(bound), + }; + // Possible optimization: if the entire node needs to be split off, + // do not move the contents but swap the nodes instead. + next_edge.move_suffix(&mut dst_node); + match (next_edge.force(), dst_node.force()) { + (Internal(next_edge), Internal(internal_dst_node)) => { + self = next_edge.descend(); + dst_node = internal_dst_node.first_edge().descend(); + bound = next_bound; + } + (Leaf(next_edge), Leaf(_)) => { + return match border { + Left(_) => next_edge.into_node().pop_last_kv(top_height), + Right(_) => None, + }; + } + _ => unreachable!(), + } + } + } +} + +impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Leaf> { + /// Backs up from a leaf node to the last KV within a subtree, and pops it. + fn pop_last_kv(self, up_to_height: usize) -> Option<(K, V)> { + let mut node = self.forget_type(); + while node.len() == 0 { + if node.height() == up_to_height { + return None; + } + node = node.ascend().ok().unwrap().into_node().forget_type(); + } + let (k, v, empty_edge) = node.pop(); + if let Some(empty_edge) = empty_edge { + empty_edge.deallocate_pillar(); + } + Some((k, v)) + } }