Auto merge of #81361 - ssomers:btree_drainy_refactor_7, r=Mark-Simulacrum

BTreeMap: lightly refactor the split_off implementation

r? `@Mark-Simulacrum`
This commit is contained in:
bors 2021-02-08 23:37:06 +00:00
commit a2704448c1
2 changed files with 64 additions and 39 deletions

View File

@ -20,6 +20,14 @@
/// We might temporarily have fewer elements during methods.
pub(super) const MIN_LEN: usize = node::MIN_LEN_AFTER_SPLIT;
// A tree in a `BTreeMap` is a tree in the `node` module with addtional invariants:
// - Keys must appear in ascending order (according to the key's type).
// - If the root node is internal, it must contain at least 1 element.
// - Every non-root node contains at least MIN_LEN elements.
//
// An empty map may be represented both by the absense of a root node or by a
// root node that is an empty leaf.
/// A map based on a B-Tree.
///
/// B-Trees represent a fundamental compromise between cache-efficiency and actually minimizing
@ -1131,20 +1139,12 @@ pub fn split_off<Q: ?Sized + Ord>(&mut self, key: &Q) -> Self
let total_num = self.len();
let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty
let mut right = Self::new();
let right_root = Self::ensure_is_owned(&mut right.root);
let right_root = left_root.split_off(key);
left_root.split_off(right_root, key);
let (new_left_len, right_len) = Root::calc_split_length(total_num, &left_root, &right_root);
self.length = new_left_len;
if left_root.height() < right_root.height() {
self.length = left_root.reborrow().calc_length();
right.length = total_num - self.len();
} else {
right.length = right_root.reborrow().calc_length();
self.length = total_num - right.len();
}
right
BTreeMap { root: Some(right_root), length: right_len }
}
/// Creates an iterator that visits all elements (key-value pairs) in

View File

@ -4,46 +4,71 @@
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)
/// Calculates the length of both trees that result from splitting up
/// a given number of distinct key-value pairs.
pub fn calc_split_length(
total_num: usize,
root_a: &Root<K, V>,
root_b: &Root<K, V>,
) -> (usize, usize) {
let (length_a, length_b);
if root_a.height() < root_b.height() {
length_a = root_a.reborrow().calc_length();
length_b = total_num - length_a;
debug_assert_eq!(length_b, root_b.reborrow().calc_length());
} else {
length_b = root_b.reborrow().calc_length();
length_a = total_num - length_b;
debug_assert_eq!(length_a, root_a.reborrow().calc_length());
}
(length_a, length_b)
}
/// Split off a tree with key-value pairs at and after the given key.
/// The result is meaningful only if the tree is ordered by key,
/// and if the ordering of `Q` corresponds to that of `K`.
/// If `self` respects all `BTreeMap` tree invariants, then both
/// `self` and the returned tree will respect those invariants.
pub fn split_off<Q: ?Sized + Ord>(&mut self, key: &Q) -> Self
where
K: Borrow<Q>,
{
debug_assert!(right_root.height() == 0);
debug_assert!(right_root.len() == 0);
let left_root = self;
for _ in 0..left_root.height() {
right_root.push_internal_level();
}
let mut right_root = Root::new_pillar(left_root.height());
let mut left_node = left_root.borrow_mut();
let mut right_node = right_root.borrow_mut();
{
let mut left_node = left_root.borrow_mut();
let mut right_node = right_root.borrow_mut();
loop {
let mut split_edge = match left_node.search_node(key) {
// key is going to the right tree
Found(kv) => kv.left_edge(),
GoDown(edge) => edge,
};
loop {
let mut split_edge = match left_node.search_node(key) {
// key is going to the right tree
Found(kv) => kv.left_edge(),
GoDown(edge) => edge,
};
split_edge.move_suffix(&mut right_node);
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!(),
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();
right_root
}
/// Creates a tree consisting of empty nodes.
fn new_pillar(height: usize) -> Self {
let mut root = Root::new();
for _ in 0..height {
root.push_internal_level();
}
root
}
/// Removes empty levels on the top, but keeps an empty leaf if the entire tree is empty.