From 003516f91a80dd8d440c932352b2fd6c6dc900b7 Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Sat, 17 Oct 2020 17:00:59 +0200 Subject: [PATCH] BTreeMap: split off most code of remove and split_off --- library/alloc/src/collections/btree/map.rs | 223 +----------------- library/alloc/src/collections/btree/mod.rs | 2 + library/alloc/src/collections/btree/remove.rs | 132 +++++++++++ library/alloc/src/collections/btree/split.rs | 104 ++++++++ 4 files changed, 239 insertions(+), 222 deletions(-) create mode 100644 library/alloc/src/collections/btree/remove.rs create mode 100644 library/alloc/src/collections/btree/split.rs diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 92cbce96054b8..20c6ebd22928b 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -13,8 +13,6 @@ use super::node::{self, marker, ForceResult::*, Handle, NodeRef}; use super::search::{self, SearchResult::*}; use super::unwrap_unchecked; -use UnderflowResult::*; - mod entry; pub use entry::{Entry, OccupiedEntry, VacantEntry}; use Entry::*; @@ -1154,40 +1152,8 @@ impl BTreeMap { let mut right = Self::new(); let right_root = Self::ensure_is_owned(&mut right.root); - for _ in 0..left_root.height() { - right_root.push_internal_level(); - } - - { - let mut left_node = left_root.node_as_mut(); - let mut right_node = right_root.node_as_mut(); - - loop { - let mut split_edge = match search::search_node(left_node, key) { - // key is going to the right tree - Found(handle) => handle.left_edge(), - GoDown(handle) => handle, - }; - split_edge.move_suffix(&mut right_node); - - match (split_edge.force(), right_node.force()) { - (Internal(edge), Internal(node)) => { - left_node = edge.descend(); - right_node = node.first_edge().descend(); - } - (Leaf(_), Leaf(_)) => { - break; - } - _ => { - unreachable!(); - } - } - } - } - - left_root.fix_right_border(); - right_root.fix_left_border(); + left_root.split_off(right_root, key); if left_root.height() < right_root.height() { self.length = left_root.node_as_ref().calc_length(); @@ -2250,193 +2216,6 @@ impl BTreeMap { } } -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. - 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() { - Leaf(leaf) => { - let (old_kv, pos) = leaf.remove(); - (old_kv, pos, false) - } - 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() < node::MIN_LEN { - match handle_underfull_node(cur_node) { - AtRoot => break, - 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() } { - Leaf(leaf) => leaf, - Internal(_) => unreachable!(), - }; - pos = unsafe { Handle::new_edge(node, idx) }; - } - - 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; - } else { - cur_node = parent.forget_type(); - at_leaf = false; - } - } - 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(); - } - } - break; - } - } - } - - // 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() }; - } - - (old_kv, pos) - } -} - -impl node::Root { - /// Removes empty levels on the top, but keep an empty leaf if the entire tree is empty. - fn fix_top(&mut self) { - while self.height() > 0 && self.node_as_ref().len() == 0 { - self.pop_internal_level(); - } - } - - fn fix_right_border(&mut self) { - self.fix_top(); - - { - let mut cur_node = self.node_as_mut(); - - while let Internal(node) = cur_node.force() { - let mut last_kv = node.last_kv(); - - if last_kv.can_merge() { - cur_node = last_kv.merge().descend(); - } else { - let right_len = last_kv.reborrow().right_edge().descend().len(); - // `MINLEN + 1` to avoid readjust if merge happens on the next level. - if right_len < node::MIN_LEN + 1 { - last_kv.bulk_steal_left(node::MIN_LEN + 1 - right_len); - } - cur_node = last_kv.right_edge().descend(); - } - } - } - - self.fix_top(); - } - - /// The symmetric clone of `fix_right_border`. - fn fix_left_border(&mut self) { - self.fix_top(); - - { - let mut cur_node = self.node_as_mut(); - - while let Internal(node) = cur_node.force() { - let mut first_kv = node.first_kv(); - - if first_kv.can_merge() { - cur_node = first_kv.merge().descend(); - } else { - let left_len = first_kv.reborrow().left_edge().descend().len(); - if left_len < node::MIN_LEN + 1 { - first_kv.bulk_steal_right(node::MIN_LEN + 1 - left_len); - } - cur_node = first_kv.left_edge().descend(); - } - } - } - - self.fix_top(); - } -} - -enum UnderflowResult<'a, K, V> { - AtRoot, - Merged(Handle, K, V, marker::Internal>, marker::Edge>, bool, usize), - Stole(bool), -} - -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 AtRoot, - }; - - // 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) - } - }; - - if handle.can_merge() { - let offset = if is_left { handle.reborrow().left_edge().descend().len() + 1 } else { 0 }; - Merged(handle.merge(), is_left, offset) - } else { - if is_left { - handle.steal_left(); - } else { - handle.steal_right(); - } - Stole(is_left) - } -} - impl> Iterator for MergeIter { type Item = (K, V); diff --git a/library/alloc/src/collections/btree/mod.rs b/library/alloc/src/collections/btree/mod.rs index ecbdacda4b618..bcc50ed561587 100644 --- a/library/alloc/src/collections/btree/mod.rs +++ b/library/alloc/src/collections/btree/mod.rs @@ -2,8 +2,10 @@ mod borrow; pub mod map; mod navigate; mod node; +mod remove; mod search; pub mod set; +mod split; #[doc(hidden)] trait Recover { diff --git a/library/alloc/src/collections/btree/remove.rs b/library/alloc/src/collections/btree/remove.rs new file mode 100644 index 0000000000000..9733b7d608425 --- /dev/null +++ b/library/alloc/src/collections/btree/remove.rs @@ -0,0 +1,132 @@ +use super::node::{self, marker, ForceResult, Handle, 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. + 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() < node::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) }; + } + + 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; + } else { + cur_node = parent.forget_type(); + at_leaf = false; + } + } + 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(); + } + } + break; + } + } + } + + // 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() }; + } + + (old_kv, pos) + } +} + +enum UnderflowResult<'a, K, V> { + AtRoot, + Merged(Handle, K, V, marker::Internal>, marker::Edge>, bool, usize), + Stole(bool), +} + +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, + }; + + // 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) + } + }; + + 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(); + } + UnderflowResult::Stole(is_left) + } +} diff --git a/library/alloc/src/collections/btree/split.rs b/library/alloc/src/collections/btree/split.rs new file mode 100644 index 0000000000000..0e6e213f6e87d --- /dev/null +++ b/library/alloc/src/collections/btree/split.rs @@ -0,0 +1,104 @@ +use super::node::{self, ForceResult::*, Root}; +use super::search::{self, SearchResult::*}; +use core::borrow::Borrow; + +impl Root { + pub fn split_off(&mut self, right_root: &mut Self, key: &Q) + where + K: Borrow, + { + debug_assert!(right_root.height() == 0); + debug_assert!(right_root.node_as_ref().len() == 0); + + let left_root = self; + for _ in 0..left_root.height() { + right_root.push_internal_level(); + } + + { + let mut left_node = left_root.node_as_mut(); + let mut right_node = right_root.node_as_mut(); + + loop { + let mut split_edge = match search::search_node(left_node, key) { + // key is going to the right tree + Found(handle) => handle.left_edge(), + GoDown(handle) => handle, + }; + + split_edge.move_suffix(&mut right_node); + + match (split_edge.force(), right_node.force()) { + (Internal(edge), Internal(node)) => { + left_node = edge.descend(); + right_node = node.first_edge().descend(); + } + (Leaf(_), Leaf(_)) => { + break; + } + _ => unreachable!(), + } + } + } + + left_root.fix_right_border(); + right_root.fix_left_border(); + } + + /// Removes empty levels on the top, but keeps an empty leaf if the entire tree is empty. + fn fix_top(&mut self) { + while self.height() > 0 && self.node_as_ref().len() == 0 { + self.pop_internal_level(); + } + } + + fn fix_right_border(&mut self) { + self.fix_top(); + + { + let mut cur_node = self.node_as_mut(); + + while let Internal(node) = cur_node.force() { + let mut last_kv = node.last_kv(); + + if last_kv.can_merge() { + cur_node = last_kv.merge().descend(); + } else { + let right_len = last_kv.reborrow().right_edge().descend().len(); + // `MINLEN + 1` to avoid readjust if merge happens on the next level. + if right_len < node::MIN_LEN + 1 { + last_kv.bulk_steal_left(node::MIN_LEN + 1 - right_len); + } + cur_node = last_kv.right_edge().descend(); + } + } + } + + self.fix_top(); + } + + /// The symmetric clone of `fix_right_border`. + fn fix_left_border(&mut self) { + self.fix_top(); + + { + let mut cur_node = self.node_as_mut(); + + while let Internal(node) = cur_node.force() { + let mut first_kv = node.first_kv(); + + if first_kv.can_merge() { + cur_node = first_kv.merge().descend(); + } else { + let left_len = first_kv.reborrow().left_edge().descend().len(); + if left_len < node::MIN_LEN + 1 { + first_kv.bulk_steal_right(node::MIN_LEN + 1 - left_len); + } + cur_node = first_kv.left_edge().descend(); + } + } + } + + self.fix_top(); + } +}