Skip to content

Commit 93e0aed

Browse files
borsbors
committed
Auto merge of #81090 - ssomers:btree_drainy_refactor_2, r=Mark-Simulacrum
BTreeMap: offer merge in variants with more clarity r? `@Mark-Simulacrum`
2 parents c4df63f + bb61cc4 commit 93e0aed

File tree

3 files changed

+60
-34
lines changed

3 files changed

+60
-34
lines changed

‎library/alloc/src/collections/btree/node.rs‎

Lines changed: 52 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1282,25 +1282,25 @@ impl<'a, K, V> BalancingContext<'a, K, V> {
12821282
self.right_child
12831283
}
12841284

1285-
/// Returns `true` if it is valid to call `.merge()` in the balancing context,
1286-
/// i.e., whether there is enough room in a node to hold the combination of
1287-
/// both adjacent child nodes, along with the key-value pair in the parent.
1285+
/// Returns whether merging is possible, i.e., whether there is enough room
1286+
/// in a node to combine the central KV with both adjacent child nodes.
12881287
pub fn can_merge(&self) -> bool {
12891288
self.left_child.len() + 1 + self.right_child.len() <= CAPACITY
12901289
}
12911290
}
12921291

12931292
impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
1294-
/// Merges the parent's key-value pair and both adjacent child nodes into
1295-
/// the left node and returns an edge handle in that expanded left node.
1296-
/// If `track_edge_idx` is given some value, the returned edge corresponds
1297-
/// to where the edge in that child node ended up,
1298-
///
1299-
/// Panics unless we `.can_merge()`.
1300-
pub fn merge(
1293+
/// Performs a merge and lets a closure decide what to return.
1294+
fn do_merge<
1295+
F: FnOnce(
1296+
NodeRef<marker::Mut<'a>, K, V, marker::Internal>,
1297+
NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
1298+
) -> R,
1299+
R,
1300+
>(
13011301
self,
1302-
track_edge_idx: Option<LeftOrRight<usize>>,
1303-
) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
1302+
result: F,
1303+
) -> R {
13041304
let Handle { node: mut parent_node, idx: parent_idx, _marker } = self.parent;
13051305
let old_parent_len = parent_node.len();
13061306
let mut left_node = self.left_child;
@@ -1310,11 +1310,6 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
13101310
let new_left_len = old_left_len + 1 + right_len;
13111311

13121312
assert!(new_left_len <= CAPACITY);
1313-
assert!(match track_edge_idx {
1314-
None => true,
1315-
Some(LeftOrRight::Left(idx)) => idx <= old_left_len,
1316-
Some(LeftOrRight::Right(idx)) => idx <= right_len,
1317-
});
13181313

13191314
unsafe {
13201315
*left_node.len_mut() = new_left_len as u16;
@@ -1353,14 +1348,47 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
13531348
} else {
13541349
Global.deallocate(right_node.node.cast(), Layout::new::<LeafNode<K, V>>());
13551350
}
1356-
1357-
let new_idx = match track_edge_idx {
1358-
None => 0,
1359-
Some(LeftOrRight::Left(idx)) => idx,
1360-
Some(LeftOrRight::Right(idx)) => old_left_len + 1 + idx,
1361-
};
1362-
Handle::new_edge(left_node, new_idx)
13631351
}
1352+
result(parent_node, left_node)
1353+
}
1354+
1355+
/// Merges the parent's key-value pair and both adjacent child nodes into
1356+
/// the left child node and returns the shrunk parent node.
1357+
///
1358+
/// Panics unless we `.can_merge()`.
1359+
pub fn merge_tracking_parent(self) -> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
1360+
self.do_merge(|parent, _child| parent)
1361+
}
1362+
1363+
/// Merges the parent's key-value pair and both adjacent child nodes into
1364+
/// the left child node and returns that child node.
1365+
///
1366+
/// Panics unless we `.can_merge()`.
1367+
pub fn merge_tracking_child(self) -> NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal> {
1368+
self.do_merge(|_parent, child| child)
1369+
}
1370+
1371+
/// Merges the parent's key-value pair and both adjacent child nodes into
1372+
/// the left child node and returns the edge handle in that child node
1373+
/// where the tracked child edge ended up,
1374+
///
1375+
/// Panics unless we `.can_merge()`.
1376+
pub fn merge_tracking_child_edge(
1377+
self,
1378+
track_edge_idx: LeftOrRight<usize>,
1379+
) -> Handle<NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>, marker::Edge> {
1380+
let old_left_len = self.left_child.len();
1381+
let right_len = self.right_child.len();
1382+
assert!(match track_edge_idx {
1383+
LeftOrRight::Left(idx) => idx <= old_left_len,
1384+
LeftOrRight::Right(idx) => idx <= right_len,
1385+
});
1386+
let child = self.merge_tracking_child();
1387+
let new_idx = match track_edge_idx {
1388+
LeftOrRight::Left(idx) => idx,
1389+
LeftOrRight::Right(idx) => old_left_len + 1 + idx,
1390+
};
1391+
unsafe { Handle::new_edge(child, new_idx) }
13641392
}
13651393

13661394
/// Removes a key-value pair from the left child and places it in the key-value storage

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

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
3333
Ok(Left(left_parent_kv)) => {
3434
debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1);
3535
if left_parent_kv.can_merge() {
36-
left_parent_kv.merge(Some(Right(idx)))
36+
left_parent_kv.merge_tracking_child_edge(Right(idx))
3737
} else {
3838
debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
3939
left_parent_kv.steal_left(idx)
@@ -42,7 +42,7 @@ impl<'a, K: 'a, V: 'a> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, mark
4242
Ok(Right(right_parent_kv)) => {
4343
debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1);
4444
if right_parent_kv.can_merge() {
45-
right_parent_kv.merge(Some(Left(idx)))
45+
right_parent_kv.merge_tracking_child_edge(Left(idx))
4646
} else {
4747
debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
4848
right_parent_kv.steal_right(idx)
@@ -124,9 +124,8 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
124124
Ok(Left(left_parent_kv)) => {
125125
debug_assert_eq!(left_parent_kv.right_child_len(), MIN_LEN - 1);
126126
if left_parent_kv.can_merge() {
127-
let pos = left_parent_kv.merge(None);
128-
let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
129-
Some(parent_edge.into_node())
127+
let parent = left_parent_kv.merge_tracking_parent();
128+
Some(parent)
130129
} else {
131130
debug_assert!(left_parent_kv.left_child_len() > MIN_LEN);
132131
left_parent_kv.steal_left(0);
@@ -136,9 +135,8 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
136135
Ok(Right(right_parent_kv)) => {
137136
debug_assert_eq!(right_parent_kv.left_child_len(), MIN_LEN - 1);
138137
if right_parent_kv.can_merge() {
139-
let pos = right_parent_kv.merge(None);
140-
let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) };
141-
Some(parent_edge.into_node())
138+
let parent = right_parent_kv.merge_tracking_parent();
139+
Some(parent)
142140
} else {
143141
debug_assert!(right_parent_kv.right_child_len() > MIN_LEN);
144142
right_parent_kv.steal_right(0);

‎library/alloc/src/collections/btree/split.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ impl<K, V> Root<K, V> {
6666
let mut last_kv = node.last_kv().consider_for_balancing();
6767

6868
if last_kv.can_merge() {
69-
cur_node = last_kv.merge(None).into_node();
69+
cur_node = last_kv.merge_tracking_child();
7070
} else {
7171
let right_len = last_kv.right_child_len();
7272
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.
@@ -93,7 +93,7 @@ impl<K, V> Root<K, V> {
9393
let mut first_kv = node.first_kv().consider_for_balancing();
9494

9595
if first_kv.can_merge() {
96-
cur_node = first_kv.merge(None).into_node();
96+
cur_node = first_kv.merge_tracking_child();
9797
} else {
9898
let left_len = first_kv.left_child_len();
9999
// `MIN_LEN + 1` to avoid readjust if merge happens on the next level.

0 commit comments

Comments
 (0)