X-Git-Url: https://git.lizzy.rs/?a=blobdiff_plain;f=library%2Falloc%2Fsrc%2Fcollections%2Fbtree%2Fremove.rs;h=c4253d4221b3b9ebdc6cc849d1fac05c456098fc;hb=4cfa5bddf105feb0c16f9023bb88ec50fcdc053f;hp=99655d3e2bf64f35c2bc0f52b1427f7e81b92fa1;hpb=b9a94c919b8cb09c186ff253360df91f223f6ef3;p=rust.git diff --git a/library/alloc/src/collections/btree/remove.rs b/library/alloc/src/collections/btree/remove.rs index 99655d3e2bf..c4253d4221b 100644 --- a/library/alloc/src/collections/btree/remove.rs +++ b/library/alloc/src/collections/btree/remove.rs @@ -1,133 +1,153 @@ use super::map::MIN_LEN; -use super::node::{marker, ForceResult, Handle, NodeRef}; +use super::node::{marker, ForceResult::*, Handle, LeftOrRight::*, NodeRef}; use super::unwrap_unchecked; use core::mem; -use core::ptr; impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInternal>, marker::KV> { - /// Removes a key/value-pair from the map, and returns that pair, as well as - /// the leaf edge corresponding to that former pair. + /// Removes a key/value-pair from the tree, and returns that pair, as well as + /// the leaf edge corresponding to that former pair. It's possible this empties + /// a root node that is internal, which the caller should pop from the map + /// holding the tree. The caller should also decrement the map's length. pub fn remove_kv_tracking( self, handle_emptied_internal_root: F, ) -> ((K, V), Handle, K, V, marker::Leaf>, marker::Edge>) { - let (old_kv, mut pos, was_internal) = match self.force() { - ForceResult::Leaf(leaf) => { - let (old_kv, pos) = leaf.remove(); - (old_kv, pos, false) - } - ForceResult::Internal(mut internal) => { - // Replace the location freed in the internal node with an - // adjacent KV, and remove that adjacent KV from its leaf. - // Always choose the adjacent KV on the left side because - // it is typically faster to pop an element from the end - // of the KV arrays without needing to shift other elements. - - let key_loc = internal.kv_mut().0 as *mut K; - let val_loc = internal.kv_mut().1 as *mut V; - - let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok(); - let to_remove = unsafe { unwrap_unchecked(to_remove) }; - - let (kv, pos) = to_remove.remove(); - - let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) }; - let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) }; - - ((old_key, old_val), pos, true) - } - }; - - // Handle underflow - let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() }; - let mut at_leaf = true; - while cur_node.len() < MIN_LEN { - match handle_underfull_node(cur_node) { - UnderflowResult::AtRoot => break, - UnderflowResult::Merged(edge, merged_with_left, offset) => { - // If we merged with our right sibling then our tracked - // position has not changed. However if we merged with our - // left sibling then our tracked position is now dangling. - if at_leaf && merged_with_left { - let idx = pos.idx() + offset; - let node = match unsafe { ptr::read(&edge).descend().force() } { - ForceResult::Leaf(leaf) => leaf, - ForceResult::Internal(_) => unreachable!(), - }; - pos = unsafe { Handle::new_edge(node, idx) }; - } + match self.force() { + Leaf(node) => node.remove_leaf_kv(handle_emptied_internal_root), + Internal(node) => node.remove_internal_kv(handle_emptied_internal_root), + } + } +} - let parent = edge.into_node(); - if parent.len() == 0 { - // The parent that was just emptied must be the root, - // because nodes on a lower level would not have been - // left with a single child. - handle_emptied_internal_root(); - break; +impl<'a, K: 'a, V: 'a> Handle, K, V, marker::Leaf>, marker::KV> { + fn remove_leaf_kv( + self, + handle_emptied_internal_root: F, + ) -> ((K, V), Handle, K, V, marker::Leaf>, marker::Edge>) { + let (old_kv, mut pos) = self.remove(); + let len = pos.reborrow().into_node().len(); + if len < MIN_LEN { + let idx = pos.idx(); + // We have to temporarily forget the child type, because there is no + // distinct node type for the immediate parents of a leaf. + let new_pos = match pos.into_node().forget_type().choose_parent_kv() { + Ok(Left(left_parent_kv)) => { + debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1); + if left_parent_kv.can_merge() { + left_parent_kv.merge(Some(Right(idx))) } else { - cur_node = parent.forget_type(); - at_leaf = false; + debug_assert!(left_parent_kv.left_child_len() > MIN_LEN); + left_parent_kv.steal_left(idx) } } - UnderflowResult::Stole(stole_from_left) => { - // Adjust the tracked position if we stole from a left sibling - if stole_from_left && at_leaf { - // SAFETY: This is safe since we just added an element to our node. - unsafe { - pos.move_next_unchecked(); - } + Ok(Right(right_parent_kv)) => { + debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1); + if right_parent_kv.can_merge() { + right_parent_kv.merge(Some(Left(idx))) + } else { + debug_assert!(right_parent_kv.right_child_len() > MIN_LEN); + right_parent_kv.steal_right(idx) } - break; } - } - } + Err(pos) => unsafe { Handle::new_edge(pos, idx) }, + }; + // SAFETY: `new_pos` is the leaf we started from or a sibling. + pos = unsafe { new_pos.cast_to_leaf_unchecked() }; - // If we deleted from an internal node then we need to compensate for - // the earlier swap and adjust the tracked position to point to the - // next element. - if was_internal { - pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() }; + // Only if we merged, the parent (if any) has shrunk, but skipping + // the following step does not pay off in benchmarks. + // + // SAFETY: We won't destroy or rearrange the leaf where `pos` is at + // by handling its parent recursively; at worst we will destroy or + // rearrange the parent through the grandparent, thus change the + // leaf's parent pointer. + if let Ok(parent) = unsafe { pos.reborrow_mut() }.into_node().ascend() { + parent.into_node().handle_shrunk_node_recursively(handle_emptied_internal_root); + } } - (old_kv, pos) } } -enum UnderflowResult<'a, K, V> { - AtRoot, - Merged(Handle, K, V, marker::Internal>, marker::Edge>, bool, usize), - Stole(bool), -} +impl<'a, K: 'a, V: 'a> Handle, K, V, marker::Internal>, marker::KV> { + fn remove_internal_kv( + self, + handle_emptied_internal_root: F, + ) -> ((K, V), Handle, K, V, marker::Leaf>, marker::Edge>) { + // Remove an adjacent KV from its leaf and then put it back in place of + // the element we were asked to remove. Prefer the left adjacent KV, + // for the reasons listed in `choose_parent_kv`. + let left_leaf_kv = self.left_edge().descend().last_leaf_edge().left_kv(); + let left_leaf_kv = unsafe { unwrap_unchecked(left_leaf_kv.ok()) }; + let (left_kv, left_hole) = left_leaf_kv.remove_leaf_kv(handle_emptied_internal_root); -fn handle_underfull_node<'a, K: 'a, V: 'a>( - node: NodeRef, K, V, marker::LeafOrInternal>, -) -> UnderflowResult<'_, K, V> { - let parent = match node.ascend() { - Ok(parent) => parent, - Err(_) => return UnderflowResult::AtRoot, - }; + // The internal node may have been stolen from or merged. Go back right + // to find where the original KV ended up. + let mut internal = unsafe { unwrap_unchecked(left_hole.next_kv().ok()) }; + let old_key = mem::replace(internal.kv_mut().0, left_kv.0); + let old_val = mem::replace(internal.kv_mut().1, left_kv.1); + let pos = internal.next_leaf_edge(); + ((old_key, old_val), pos) + } +} - // Prefer the left KV if it exists. Merging with the left side is faster, - // since merging happens towards the left and `node` has fewer elements. - // Stealing from the left side is faster, since we can pop from the end of - // the KV arrays. - let (is_left, mut handle) = match parent.left_kv() { - Ok(left) => (true, left), - Err(parent) => { - let right = unsafe { unwrap_unchecked(parent.right_kv().ok()) }; - (false, right) +impl<'a, K: 'a, V: 'a> NodeRef, K, V, marker::Internal> { + /// Stocks up a possibly underfull internal node, recursively. + /// Climbs up until it reaches an ancestor that has elements to spare or the root. + fn handle_shrunk_node_recursively(mut self, handle_emptied_internal_root: F) { + loop { + self = match self.len() { + 0 => { + // An empty node must be the root, because length is only + // reduced by one, and non-root underfull nodes are stocked up, + // so non-root nodes never have fewer than MIN_LEN - 1 elements. + debug_assert!(self.ascend().is_err()); + handle_emptied_internal_root(); + return; + } + 1..MIN_LEN => { + if let Some(parent) = self.handle_underfull_node_locally() { + parent + } else { + return; + } + } + _ => return, + } } - }; + } - if handle.can_merge() { - let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 }; - UnderflowResult::Merged(handle.merge(), is_left, offset) - } else { - if is_left { - handle.steal_left(); - } else { - handle.steal_right(); + /// Stocks up an underfull internal node, possibly at the cost of shrinking + /// its parent instead, which is then returned. + fn handle_underfull_node_locally( + self, + ) -> Option, K, V, marker::Internal>> { + match self.forget_type().choose_parent_kv() { + Ok(Left(left_parent_kv)) => { + debug_assert!(left_parent_kv.right_child_len() == MIN_LEN - 1); + if left_parent_kv.can_merge() { + let pos = left_parent_kv.merge(None); + let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) }; + Some(parent_edge.into_node()) + } else { + debug_assert!(left_parent_kv.left_child_len() > MIN_LEN); + left_parent_kv.steal_left(0); + None + } + } + Ok(Right(right_parent_kv)) => { + debug_assert!(right_parent_kv.left_child_len() == MIN_LEN - 1); + if right_parent_kv.can_merge() { + let pos = right_parent_kv.merge(None); + let parent_edge = unsafe { unwrap_unchecked(pos.into_node().ascend().ok()) }; + Some(parent_edge.into_node()) + } else { + debug_assert!(right_parent_kv.right_child_len() > MIN_LEN); + right_parent_kv.steal_right(0); + None + } + } + Err(_) => None, } - UnderflowResult::Stole(is_left) } }