diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index ecc28731874..1a62fc682e6 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -21,6 +21,14 @@ use Entry::*; /// 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 @@ -1100,20 +1108,12 @@ impl BTreeMap { 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 which uses a closure to determine if an element should be removed. diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs index 62c5e3a46d7..1921982464a 100644 --- a/library/alloc/src/collections/btree/split.rs +++ b/library/alloc/src/collections/btree/split.rs @@ -4,46 +4,71 @@ use super::search::SearchResult::*; use core::borrow::Borrow; impl Root { - pub fn split_off(&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, + root_b: &Root, + ) -> (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(&mut self, key: &Q) -> Self where K: Borrow, { - 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.