From ca253cab3689c0d525dabda10e35c469839b2c4b Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Mon, 13 Jul 2020 16:51:37 +0200 Subject: [PATCH 1/2] Clean up or comment every unwrap in BTreeMap's main code. --- src/liballoc/collections/btree/map.rs | 81 +++++++++++--------------- src/liballoc/collections/btree/node.rs | 5 ++ 2 files changed, 40 insertions(+), 46 deletions(-) diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs index f3781db1cf784..d94c9539e6dc5 100644 --- a/src/liballoc/collections/btree/map.rs +++ b/src/liballoc/collections/btree/map.rs @@ -151,7 +151,7 @@ impl Clone for BTreeMap { let mut out_tree = BTreeMap { root: Some(node::Root::new_leaf()), length: 0 }; { - let root = out_tree.root.as_mut().unwrap(); + let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped let mut out_node = match root.as_mut().force() { Leaf(leaf) => leaf, Internal(_) => unreachable!(), @@ -171,14 +171,10 @@ impl Clone for BTreeMap { } Internal(internal) => { let mut out_tree = clone_subtree(internal.first_edge().descend()); - out_tree.ensure_root_is_owned(); { - // Ideally we'd use the return of ensure_root_is_owned - // instead of re-unwrapping here but unfortunately that - // borrows all of out_tree and we need access to the - // length below. - let mut out_node = out_tree.root.as_mut().unwrap().push_level(); + let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root); + let mut out_node = out_root.push_level(); let mut in_edge = internal.first_edge(); while let Ok(kv) = in_edge.right_kv() { let (k, v) = kv.into_kv(); @@ -212,7 +208,7 @@ impl Clone for BTreeMap { // Ord` constraint, which this method lacks. BTreeMap { root: None, length: 0 } } else { - clone_subtree(self.root.as_ref().unwrap().as_ref()) + clone_subtree(self.root.as_ref().unwrap().as_ref()) // unwrap succeeds because not empty } } } @@ -243,8 +239,8 @@ where } fn replace(&mut self, key: K) -> Option { - self.ensure_root_is_owned(); - match search::search_tree::, K, (), K>(self.root.as_mut()?.as_mut(), &key) { + let root = Self::ensure_is_owned(&mut self.root); + match search::search_tree::, K, (), K>(root.as_mut(), &key) { Found(handle) => Some(mem::replace(handle.into_kv_mut().0, key)), GoDown(handle) => { VacantEntry { key, handle, length: &mut self.length, _marker: PhantomData } @@ -943,7 +939,6 @@ impl BTreeMap { // Second, we build a tree from the sorted sequence in linear time. self.from_sorted_iter(iter); - self.fix_right_edge(); } /// Constructs a double-ended iterator over a sub-range of elements in the map. @@ -1058,8 +1053,8 @@ impl BTreeMap { #[stable(feature = "rust1", since = "1.0.0")] pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { // FIXME(@porglezomp) Avoid allocating if we don't insert - self.ensure_root_is_owned(); - match search::search_tree(self.root.as_mut().unwrap().as_mut(), &key) { + let root = Self::ensure_is_owned(&mut self.root); + match search::search_tree(root.as_mut(), &key) { Found(handle) => { Occupied(OccupiedEntry { handle, length: &mut self.length, _marker: PhantomData }) } @@ -1070,8 +1065,8 @@ impl BTreeMap { } fn from_sorted_iter>(&mut self, iter: I) { - self.ensure_root_is_owned(); - let mut cur_node = self.root.as_mut().unwrap().as_mut().last_leaf_edge().into_node(); + let root = Self::ensure_is_owned(&mut self.root); + let mut cur_node = root.as_mut().last_leaf_edge().into_node(); // Iterate through all key-value pairs, pushing them into nodes at the right level. for (key, value) in iter { // Try to push key-value pair into the current leaf node. @@ -1116,11 +1111,12 @@ impl BTreeMap { self.length += 1; } + Self::fix_right_edge(root) } - fn fix_right_edge(&mut self) { + fn fix_right_edge(root: &mut node::Root) { // Handle underfull nodes, start from the top. - let mut cur_node = self.root.as_mut().unwrap().as_mut(); + let mut cur_node = root.as_mut(); while let Internal(internal) = cur_node.force() { // Check if right-most child is underfull. let mut last_edge = internal.last_edge(); @@ -1179,16 +1175,17 @@ impl BTreeMap { } let total_num = self.len(); + let left_root = self.root.as_mut().unwrap(); // unwrap succeeds because not empty let mut right = Self::new(); - let right_root = right.ensure_root_is_owned(); - for _ in 0..(self.root.as_ref().unwrap().as_ref().height()) { + let right_root = Self::ensure_is_owned(&mut right.root); + for _ in 0..left_root.height() { right_root.push_level(); } { - let mut left_node = self.root.as_mut().unwrap().as_mut(); - let mut right_node = right.root.as_mut().unwrap().as_mut(); + let mut left_node = left_root.as_mut(); + let mut right_node = right_root.as_mut(); loop { let mut split_edge = match search::search_node(left_node, key) { @@ -1214,12 +1211,10 @@ impl BTreeMap { } } - self.fix_right_border(); - right.fix_left_border(); + Self::fix_right_border(left_root); + Self::fix_left_border(right_root); - if self.root.as_ref().unwrap().as_ref().height() - < right.root.as_ref().unwrap().as_ref().height() - { + if left_root.height() < right_root.height() { self.recalc_length(); right.length = total_num - self.len(); } else { @@ -1303,23 +1298,17 @@ impl BTreeMap { } /// Removes empty levels on the top. - fn fix_top(&mut self) { - loop { - { - let node = self.root.as_ref().unwrap().as_ref(); - if node.height() == 0 || node.len() > 0 { - break; - } - } - self.root.as_mut().unwrap().pop_level(); + fn fix_top(root: &mut node::Root) { + while root.height() > 0 && root.as_ref().len() == 0 { + root.pop_level(); } } - fn fix_right_border(&mut self) { - self.fix_top(); + fn fix_right_border(root: &mut node::Root) { + Self::fix_top(root); { - let mut cur_node = self.root.as_mut().unwrap().as_mut(); + let mut cur_node = root.as_mut(); while let Internal(node) = cur_node.force() { let mut last_kv = node.last_kv(); @@ -1337,15 +1326,15 @@ impl BTreeMap { } } - self.fix_top(); + Self::fix_top(root); } /// The symmetric clone of `fix_right_border`. - fn fix_left_border(&mut self) { - self.fix_top(); + fn fix_left_border(root: &mut node::Root) { + Self::fix_top(root); { - let mut cur_node = self.root.as_mut().unwrap().as_mut(); + let mut cur_node = root.as_mut(); while let Internal(node) = cur_node.force() { let mut first_kv = node.first_kv(); @@ -1362,7 +1351,7 @@ impl BTreeMap { } } - self.fix_top(); + Self::fix_top(root); } } @@ -2321,9 +2310,9 @@ impl BTreeMap { } /// If the root node is the empty (non-allocated) root node, allocate our - /// own node. - fn ensure_root_is_owned(&mut self) -> &mut node::Root { - self.root.get_or_insert_with(node::Root::new_leaf) + /// own node. Is an associated function to avoid borrowing the entire BTreeMap. + fn ensure_is_owned(root: &mut Option>) -> &mut node::Root { + root.get_or_insert_with(node::Root::new_leaf) } } diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs index ce74d4f8ee688..f7bd64608d63c 100644 --- a/src/liballoc/collections/btree/node.rs +++ b/src/liballoc/collections/btree/node.rs @@ -153,6 +153,11 @@ unsafe impl Sync for Root {} unsafe impl Send for Root {} impl Root { + /// Returns the number of levels below the root. + pub fn height(&self) -> usize { + self.height + } + /// Returns a new owned tree, with its own root node that is initially empty. pub fn new_leaf() -> Self { Root { node: BoxedNode::from_leaf(Box::new(unsafe { LeafNode::new() })), height: 0 } From b82d332c52dde1680b21c0281f08cc5f30edc082 Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Tue, 14 Jul 2020 13:23:15 +0200 Subject: [PATCH 2/2] Separate off BTreeMap support functions and loose their irrelevant bounds --- src/liballoc/collections/btree/map.rs | 120 +++++++++++++------------- 1 file changed, 61 insertions(+), 59 deletions(-) diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs index d94c9539e6dc5..bf5748739d470 100644 --- a/src/liballoc/collections/btree/map.rs +++ b/src/liballoc/collections/btree/map.rs @@ -1211,8 +1211,8 @@ impl BTreeMap { } } - Self::fix_right_border(left_root); - Self::fix_left_border(right_root); + left_root.fix_right_border(); + right_root.fix_left_border(); if left_root.height() < right_root.height() { self.recalc_length(); @@ -1296,63 +1296,6 @@ impl BTreeMap { self.length = dfs(self.root.as_ref().unwrap().as_ref()); } - - /// Removes empty levels on the top. - fn fix_top(root: &mut node::Root) { - while root.height() > 0 && root.as_ref().len() == 0 { - root.pop_level(); - } - } - - fn fix_right_border(root: &mut node::Root) { - Self::fix_top(root); - - { - let mut cur_node = root.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(root); - } - - /// The symmetric clone of `fix_right_border`. - fn fix_left_border(root: &mut node::Root) { - Self::fix_top(root); - - { - let mut cur_node = root.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(root); - } } #[stable(feature = "rust1", since = "1.0.0")] @@ -2814,6 +2757,65 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInter } } +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.as_ref().len() == 0 { + self.pop_level(); + } + } + + fn fix_right_border(&mut self) { + self.fix_top(); + + { + let mut cur_node = self.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.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),