Skip to content

Commit bb61cc4

Browse files
ssomersssomers
committed
BTreeMap: offer merge in variants with more clarity
1 parent efdb859 commit bb61cc4

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
@@ -1276,25 +1276,25 @@ impl<'a, K, V> BalancingContext<'a, K, V> {
12761276
self.right_child
12771277
}
12781278

1279-
/// Returns `true` if it is valid to call `.merge()` in the balancing context,
1280-
/// i.e., whether there is enough room in a node to hold the combination of
1281-
/// both adjacent child nodes, along with the key-value pair in the parent.
1279+
/// Returns whether merging is possible, i.e., whether there is enough room
1280+
/// in a node to combine the central KV with both adjacent child nodes.
12821281
pub fn can_merge(&self) -> bool {
12831282
self.left_child.len() + 1 + self.right_child.len() <= CAPACITY
12841283
}
12851284
}
12861285

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

13061306
assert!(new_left_len <= CAPACITY);
1307-
assert!(match track_edge_idx {
1308-
None => true,
1309-
Some(LeftOrRight::Left(idx)) => idx <= old_left_len,
1310-
Some(LeftOrRight::Right(idx)) => idx <= right_len,
1311-
});
13121307

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

13601388
/// 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)