Skip to content

Commit

Permalink
Auto merge of #74827 - ssomers:btree_cleanup_insert, r=Mark-Simulacrum
Browse files Browse the repository at this point in the history
Move bulk of BTreeMap::insert method down to new method on handle

Adjust the boundary between the map and node layers for insertion: do more in the node layer, keep root manipulation and pointer dereferencing separate. No change in undefined behaviour or performance.

r? @Mark-Simulacrum
  • Loading branch information
bors committed Aug 3, 2020
2 parents c186aed + f5c47fa commit 829d69b
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 37 deletions.
41 changes: 9 additions & 32 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2465,40 +2465,17 @@ impl<'a, K: Ord, V> VacantEntry<'a, K, V> {
pub fn insert(self, value: V) -> &'a mut V {
*self.length += 1;

let out_ptr;

let mut ins_k;
let mut ins_v;
let mut ins_edge;

let mut cur_parent = match self.handle.insert(self.key, value) {
(Fit(handle), _) => return handle.into_kv_mut().1,
(Split(left, k, v, right), ptr) => {
ins_k = k;
ins_v = v;
ins_edge = right;
out_ptr = ptr;
left.ascend().map_err(|n| n.into_root_mut())
let out_ptr = match self.handle.insert_recursing(self.key, value) {
(Fit(_), val_ptr) => val_ptr,
(Split(ins), val_ptr) => {
let root = ins.left.into_root_mut();
root.push_internal_level().push(ins.k, ins.v, ins.right);
val_ptr
}
};

loop {
match cur_parent {
Ok(parent) => match parent.insert(ins_k, ins_v, ins_edge) {
Fit(_) => return unsafe { &mut *out_ptr },
Split(left, k, v, right) => {
ins_k = k;
ins_v = v;
ins_edge = right;
cur_parent = left.ascend().map_err(|n| n.into_root_mut());
}
},
Err(root) => {
root.push_internal_level().push(ins_k, ins_v, ins_edge);
return unsafe { &mut *out_ptr };
}
}
}
// Now that we have finished growing the tree using borrowed references,
// dereference the pointer to a part of it, that we picked up along the way.
unsafe { &mut *out_ptr }
}
}

Expand Down
66 changes: 61 additions & 5 deletions library/alloc/src/collections/btree/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -843,7 +843,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
/// this edge. This method splits the node if there isn't enough room.
///
/// The returned pointer points to the inserted value.
pub fn insert(mut self, key: K, val: V) -> (InsertResult<'a, K, V, marker::Leaf>, *mut V) {
fn insert(mut self, key: K, val: V) -> (InsertResult<'a, K, V, marker::Leaf>, *mut V) {
if self.node.len() < CAPACITY {
let ptr = self.insert_fit(key, val);
let kv = unsafe { Handle::new_kv(self.node, self.idx) };
Expand All @@ -862,7 +862,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
.insert_fit(key, val)
}
};
(InsertResult::Split(left, k, v, right), ptr)
(InsertResult::Split(SplitResult { left: left.forget_type(), k, v, right }), ptr)
}
}
}
Expand Down Expand Up @@ -918,7 +918,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
/// Inserts a new key/value pair and an edge that will go to the right of that new pair
/// between this edge and the key/value pair to the right of this edge. This method splits
/// the node if there isn't enough room.
pub fn insert(
fn insert(
mut self,
key: K,
val: V,
Expand Down Expand Up @@ -946,7 +946,43 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Internal>, marker::
.insert_fit(key, val, edge);
}
}
InsertResult::Split(left, k, v, right)
InsertResult::Split(SplitResult { left: left.forget_type(), k, v, right })
}
}
}

impl<'a, K: 'a, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
/// Inserts a new key/value pair between the key/value pairs to the right and left of
/// this edge. This method splits the node if there isn't enough room, and tries to
/// insert the split off portion into the parent node recursively, until the root is reached.
///
/// If the returned result is a `Fit`, its handle's node can be this edge's node or an ancestor.
/// If the returned result is a `Split`, the `left` field will be the root node.
/// The returned pointer points to the inserted value.
pub fn insert_recursing(
self,
key: K,
value: V,
) -> (InsertResult<'a, K, V, marker::LeafOrInternal>, *mut V) {
let (mut split, val_ptr) = match self.insert(key, value) {
(InsertResult::Fit(handle), ptr) => {
return (InsertResult::Fit(handle.forget_node_type()), ptr);
}
(InsertResult::Split(split), val_ptr) => (split, val_ptr),
};

loop {
split = match split.left.ascend() {
Ok(parent) => match parent.insert(split.k, split.v, split.right) {
InsertResult::Fit(handle) => {
return (InsertResult::Fit(handle.forget_node_type()), val_ptr);
}
InsertResult::Split(split) => split,
},
Err(root) => {
return (InsertResult::Split(SplitResult { left: root, ..split }), val_ptr);
}
};
}
}
}
Expand Down Expand Up @@ -1389,6 +1425,14 @@ impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::K
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
pub fn forget_node_type(
self,
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
unsafe { Handle::new_kv(self.node.forget_type(), self.idx) }
}
}

impl<BorrowType, K, V, HandleType>
Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, HandleType>
{
Expand Down Expand Up @@ -1455,9 +1499,21 @@ pub enum ForceResult<Leaf, Internal> {
Internal(Internal),
}

/// Result of insertion, when a node needed to expand beyond its capacity.
/// Does not distinguish between `Leaf` and `Internal` because `Root` doesn't.
pub struct SplitResult<'a, K, V> {
// Altered node in existing tree with elements and edges that belong to the left of `k`.
pub left: NodeRef<marker::Mut<'a>, K, V, marker::LeafOrInternal>,
// Some key and value split off, to be inserted elsewhere.
pub k: K,
pub v: V,
// Owned, unattached, new node with elements and edges that belong to the right of `k`.
pub right: Root<K, V>,
}

pub enum InsertResult<'a, K, V, Type> {
Fit(Handle<NodeRef<marker::Mut<'a>, K, V, Type>, marker::KV>),
Split(NodeRef<marker::Mut<'a>, K, V, Type>, K, V, Root<K, V>),
Split(SplitResult<'a, K, V>),
}

pub mod marker {
Expand Down

0 comments on commit 829d69b

Please sign in to comment.