Skip to content

Commit

Permalink
BTree: add split_off_range methods
Browse files Browse the repository at this point in the history
  • Loading branch information
ssomers committed Jun 17, 2022
1 parent d40f24e commit b7faf86
Show file tree
Hide file tree
Showing 8 changed files with 1,085 additions and 51 deletions.
333 changes: 322 additions & 11 deletions library/alloc/src/collections/btree/fix.rs

Large diffs are not rendered by default.

62 changes: 62 additions & 0 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1307,6 +1307,68 @@ impl<K, V, A: Allocator> BTreeMap<K, V, A> {
}
}

/// 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_split_off_range)]
/// 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_split_off_range", issue = "81074")]
pub fn split_off_range<T: ?Sized, R>(&mut self, range: R) -> Self
where
T: Ord,
K: Borrow<T> + Ord,
R: RangeBounds<T>,
A: Clone,
{
let alloc = (*self.alloc).clone();
if self.is_empty() {
return Self::new_in(alloc);
}

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, &alloc);
right_root.fix_both_borders(&alloc);

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, alloc: ManuallyDrop::new(alloc) }
}

/// 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
Expand Down
28 changes: 28 additions & 0 deletions library/alloc/src/collections/btree/map/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

// Minimum number of elements to insert, to guarantee a tree with 2 levels,
// i.e., a tree who's root is an internal node at height 1, with edges to leaf nodes.
// It's not the minimum size: removing an element from such a tree does not always reduce height.
Expand All @@ -27,6 +29,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<Item = &'a mut T>) {
Expand Down Expand Up @@ -174,6 +182,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.
Expand Down
Loading

0 comments on commit b7faf86

Please sign in to comment.