Skip to content

Commit cbbdb44

Browse files
committed
BTreeMap: gather and decompose reusable tree fixing functions
1 parent f4008fe commit cbbdb44

File tree

5 files changed

+176
-146
lines changed

5 files changed

+176
-146
lines changed

library/alloc/src/collections/btree/append.rs

+1-22
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
1-
use super::map::MIN_LEN;
21
use super::merge_iter::MergeIterInner;
3-
use super::node::{self, ForceResult::*, Root};
2+
use super::node::{self, Root};
43
use core::iter::FusedIterator;
54

65
impl<K, V> Root<K, V> {
@@ -83,26 +82,6 @@ impl<K, V> Root<K, V> {
8382
}
8483
self.fix_right_border_of_plentiful();
8584
}
86-
87-
/// Stock up any underfull nodes on the right border of the tree.
88-
/// The other nodes, those that are not the root nor a rightmost edge,
89-
/// must have MIN_LEN elements to spare.
90-
fn fix_right_border_of_plentiful(&mut self) {
91-
let mut cur_node = self.borrow_mut();
92-
while let Internal(internal) = cur_node.force() {
93-
// Check if right-most child is underfull.
94-
let mut last_kv = internal.last_kv().consider_for_balancing();
95-
debug_assert!(last_kv.left_child_len() >= MIN_LEN * 2);
96-
let right_child_len = last_kv.right_child_len();
97-
if right_child_len < MIN_LEN {
98-
// We need to steal.
99-
last_kv.bulk_steal_left(MIN_LEN - right_child_len);
100-
}
101-
102-
// Go further down.
103-
cur_node = last_kv.into_right_child();
104-
}
105-
}
10685
}
10786

10887
// An iterator for merging two sorted sequences into one
+171
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
use super::map::MIN_LEN;
2+
use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef, Root};
3+
4+
impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
5+
/// Stocks up a possibly underfull node by merging with or stealing from a
6+
/// sibling. If succesful but at the cost of shrinking the parent node,
7+
/// returns that shrunk parent node. Returns an `Err` if the node is
8+
/// an empty root.
9+
fn fix_node_through_parent(
10+
self,
11+
) -> Result<Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>>, Self> {
12+
let len = self.len();
13+
if len >= MIN_LEN {
14+
Ok(None)
15+
} else {
16+
match self.choose_parent_kv() {
17+
Ok(Left(mut left_parent_kv)) => {
18+
if left_parent_kv.can_merge() {
19+
let parent = left_parent_kv.merge_tracking_parent();
20+
Ok(Some(parent))
21+
} else {
22+
left_parent_kv.bulk_steal_left(MIN_LEN - len);
23+
Ok(None)
24+
}
25+
}
26+
Ok(Right(mut right_parent_kv)) => {
27+
if right_parent_kv.can_merge() {
28+
let parent = right_parent_kv.merge_tracking_parent();
29+
Ok(Some(parent))
30+
} else {
31+
right_parent_kv.bulk_steal_right(MIN_LEN - len);
32+
Ok(None)
33+
}
34+
}
35+
Err(root) => {
36+
if len > 0 {
37+
Ok(None)
38+
} else {
39+
Err(root)
40+
}
41+
}
42+
}
43+
}
44+
}
45+
}
46+
47+
impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
48+
/// Stocks up a possibly underfull node, and if that causes its parent node
49+
/// to shrink, stocks up the parent, recursively.
50+
/// Returns `true` if it fixed the tree, `false` if it couldn't because the
51+
/// root node became empty.
52+
///
53+
/// This method does not expect ancestors to already be underfull upon entry
54+
/// and panics if it encounters an empty ancestor.
55+
pub fn fix_node_and_affected_ancestors(mut self) -> bool {
56+
loop {
57+
match self.fix_node_through_parent() {
58+
Ok(Some(parent)) => self = parent.forget_type(),
59+
Ok(None) => return true,
60+
Err(_) => return false,
61+
}
62+
}
63+
}
64+
}
65+
66+
impl<K, V> Root<K, V> {
67+
/// Removes empty levels on the top, but keeps an empty leaf if the entire tree is empty.
68+
pub fn fix_top(&mut self) {
69+
while self.height() > 0 && self.len() == 0 {
70+
self.pop_internal_level();
71+
}
72+
}
73+
74+
/// Stocks up or merge away any underfull nodes on the right border of the
75+
/// tree. The other nodes, those that are not the root nor a rightmost edge,
76+
/// must already have at least MIN_LEN elements.
77+
pub fn fix_right_border(&mut self) {
78+
self.fix_top();
79+
if self.len() > 0 {
80+
self.borrow_mut().last_kv().fix_right_border_of_right_edge();
81+
self.fix_top();
82+
}
83+
}
84+
85+
/// The symmetric clone of `fix_right_border`.
86+
pub fn fix_left_border(&mut self) {
87+
self.fix_top();
88+
if self.len() > 0 {
89+
self.borrow_mut().first_kv().fix_left_border_of_left_edge();
90+
self.fix_top();
91+
}
92+
}
93+
94+
/// Stock up any underfull nodes on the right border of the tree.
95+
/// The other nodes, those that are not the root nor a rightmost edge,
96+
/// must be prepared to have up to MIN_LEN elements stolen.
97+
pub fn fix_right_border_of_plentiful(&mut self) {
98+
let mut cur_node = self.borrow_mut();
99+
while let Internal(internal) = cur_node.force() {
100+
// Check if right-most child is underfull.
101+
let mut last_kv = internal.last_kv().consider_for_balancing();
102+
debug_assert!(last_kv.left_child_len() >= MIN_LEN * 2);
103+
let right_child_len = last_kv.right_child_len();
104+
if right_child_len < MIN_LEN {
105+
// We need to steal.
106+
last_kv.bulk_steal_left(MIN_LEN - right_child_len);
107+
}
108+
109+
// Go further down.
110+
cur_node = last_kv.into_right_child();
111+
}
112+
}
113+
}
114+
115+
impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::KV> {
116+
fn fix_left_border_of_left_edge(mut self) {
117+
while let Internal(internal_kv) = self.force() {
118+
self = internal_kv.fix_left_child().first_kv();
119+
debug_assert!(self.reborrow().into_node().len() > MIN_LEN);
120+
}
121+
}
122+
123+
fn fix_right_border_of_right_edge(mut self) {
124+
while let Internal(internal_kv) = self.force() {
125+
self = internal_kv.fix_right_child().last_kv();
126+
debug_assert!(self.reborrow().into_node().len() > MIN_LEN);
127+
}
128+
}
129+
}
130+
131+
impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::KV> {
132+
/// Stocks up the left child, assuming the right child isn't underfull, and
133+
/// provisions an extra element to allow merging its children in turn
134+
/// without becoming underfull.
135+
/// Returns the left child.
136+
fn fix_left_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
137+
let mut internal_kv = self.consider_for_balancing();
138+
let left_len = internal_kv.left_child_len();
139+
debug_assert!(internal_kv.right_child_len() >= MIN_LEN);
140+
if internal_kv.can_merge() {
141+
internal_kv.merge_tracking_child()
142+
} else {
143+
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
144+
let count = (MIN_LEN + 1).saturating_sub(left_len);
145+
if count > 0 {
146+
internal_kv.bulk_steal_right(count);
147+
}
148+
internal_kv.into_left_child()
149+
}
150+
}
151+
152+
/// Stocks up the right child, assuming the left child isn't underfull, and
153+
/// provisions an extra element to allow merging its children in turn
154+
/// without becoming underfull.
155+
/// Returns wherever the right child ended up.
156+
fn fix_right_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
157+
let mut internal_kv = self.consider_for_balancing();
158+
let right_len = internal_kv.right_child_len();
159+
debug_assert!(internal_kv.left_child_len() >= MIN_LEN);
160+
if internal_kv.can_merge() {
161+
internal_kv.merge_tracking_child()
162+
} else {
163+
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
164+
let count = (MIN_LEN + 1).saturating_sub(right_len);
165+
if count > 0 {
166+
internal_kv.bulk_steal_left(count);
167+
}
168+
internal_kv.into_right_child()
169+
}
170+
}
171+
}

library/alloc/src/collections/btree/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
mod append;
22
mod borrow;
3+
mod fix;
34
pub mod map;
45
mod mem;
56
mod merge_iter;

library/alloc/src/collections/btree/remove.rs

+3-60
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,9 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
6060
// rearrange the parent through the grandparent, thus change the
6161
// link to the parent inside the leaf.
6262
if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() {
63-
parent.into_node().handle_shrunk_node_recursively(handle_emptied_internal_root);
63+
if !parent.into_node().forget_type().fix_node_and_affected_ancestors() {
64+
handle_emptied_internal_root();
65+
}
6466
}
6567
}
6668
(old_kv, pos)
@@ -87,62 +89,3 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
8789
(old_kv, pos)
8890
}
8991
}
90-
91-
impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
92-
/// Stocks up a possibly underfull internal node and its ancestors,
93-
/// until it reaches an ancestor that has elements to spare or is the root.
94-
fn handle_shrunk_node_recursively<F: FnOnce()>(mut self, handle_emptied_internal_root: F) {
95-
loop {
96-
self = match self.len() {
97-
0 => {
98-
// An empty node must be the root, because length is only
99-
// reduced by one, and non-root underfull nodes are stocked up,
100-
// so non-root nodes never have fewer than MIN_LEN - 1 elements.
101-
debug_assert!(self.ascend().is_err());
102-
handle_emptied_internal_root();
103-
return;
104-
}
105-
1..MIN_LEN => {
106-
if let Some(parent) = self.handle_underfull_node_locally() {
107-
parent
108-
} else {
109-
return;
110-
}
111-
}
112-
_ => return,
113-
}
114-
}
115-
}
116-
117-
/// Stocks up an underfull internal node, possibly at the cost of shrinking
118-
/// its parent instead, which is then returned.
119-
fn handle_underfull_node_locally(
120-
self,
121-
) -> Option<NodeRef<marker::Mut<'a>, K, V, marker::Internal>> {
122-
match self.forget_type().choose_parent_kv() {
123-
Ok(Left(mut left_parent_kv)) => {
124-
debug_assert_eq!(left_parent_kv.right_child_len(), MIN_LEN - 1);
125-
if left_parent_kv.can_merge() {
126-
let parent = left_parent_kv.merge_tracking_parent();
127-
Some(parent)
128-
} else {
129-
debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
130-
left_parent_kv.bulk_steal_left(1);
131-
None
132-
}
133-
}
134-
Ok(Right(mut right_parent_kv)) => {
135-
debug_assert_eq!(right_parent_kv.left_child_len(), MIN_LEN - 1);
136-
if right_parent_kv.can_merge() {
137-
let parent = right_parent_kv.merge_tracking_parent();
138-
Some(parent)
139-
} else {
140-
debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
141-
right_parent_kv.bulk_steal_right(1);
142-
None
143-
}
144-
}
145-
Err(_) => None,
146-
}
147-
}
148-
}
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
use super::map::MIN_LEN;
21
use super::node::{ForceResult::*, Root};
32
use super::search::SearchResult::*;
43
use core::borrow::Borrow;
@@ -70,67 +69,4 @@ impl<K, V> Root<K, V> {
7069
}
7170
root
7271
}
73-
74-
/// Removes empty levels on the top, but keeps an empty leaf if the entire tree is empty.
75-
fn fix_top(&mut self) {
76-
while self.height() > 0 && self.len() == 0 {
77-
self.pop_internal_level();
78-
}
79-
}
80-
81-
/// Stock up or merge away any underfull nodes on the right border of the
82-
/// tree. The other nodes, those that are not the root nor a rightmost edge,
83-
/// must already have at least MIN_LEN elements.
84-
fn fix_right_border(&mut self) {
85-
self.fix_top();
86-
87-
{
88-
let mut cur_node = self.borrow_mut();
89-
90-
while let Internal(node) = cur_node.force() {
91-
let mut last_kv = node.last_kv().consider_for_balancing();
92-
93-
if last_kv.can_merge() {
94-
cur_node = last_kv.merge_tracking_child();
95-
} else {
96-
let right_len = last_kv.right_child_len();
97-
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
98-
if right_len < MIN_LEN + 1 {
99-
last_kv.bulk_steal_left(MIN_LEN + 1 - right_len);
100-
}
101-
cur_node = last_kv.into_right_child();
102-
}
103-
debug_assert!(cur_node.len() > MIN_LEN);
104-
}
105-
}
106-
107-
self.fix_top();
108-
}
109-
110-
/// The symmetric clone of `fix_right_border`.
111-
fn fix_left_border(&mut self) {
112-
self.fix_top();
113-
114-
{
115-
let mut cur_node = self.borrow_mut();
116-
117-
while let Internal(node) = cur_node.force() {
118-
let mut first_kv = node.first_kv().consider_for_balancing();
119-
120-
if first_kv.can_merge() {
121-
cur_node = first_kv.merge_tracking_child();
122-
} else {
123-
let left_len = first_kv.left_child_len();
124-
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
125-
if left_len < MIN_LEN + 1 {
126-
first_kv.bulk_steal_right(MIN_LEN + 1 - left_len);
127-
}
128-
cur_node = first_kv.into_left_child();
129-
}
130-
debug_assert!(cur_node.len() > MIN_LEN);
131-
}
132-
}
133-
134-
self.fix_top();
135-
}
13672
}

0 commit comments

Comments
 (0)