diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 6184051316e..1db629c3bdf 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -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 } } } diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index b767d9ebed7..873713302c2 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -843,7 +843,7 @@ impl<'a, K, V> Handle, 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) }; @@ -862,7 +862,7 @@ impl<'a, K, V> Handle, 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) } } } @@ -918,7 +918,7 @@ impl<'a, K, V> Handle, 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, @@ -946,7 +946,43 @@ impl<'a, K, V> Handle, 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, 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); + } + }; } } } @@ -1389,6 +1425,14 @@ impl Handle, marker::K } } +impl Handle, marker::KV> { + pub fn forget_node_type( + self, + ) -> Handle, marker::KV> { + unsafe { Handle::new_kv(self.node.forget_type(), self.idx) } + } +} + impl Handle, HandleType> { @@ -1455,9 +1499,21 @@ pub enum ForceResult { 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, 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, +} + pub enum InsertResult<'a, K, V, Type> { Fit(Handle, K, V, Type>, marker::KV>), - Split(NodeRef, K, V, Type>, K, V, Root), + Split(SplitResult<'a, K, V>), } pub mod marker {