From 85a78793412c7cf6bdde982f7de2fc3de204817c Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Wed, 5 Aug 2020 10:36:12 +0200 Subject: [PATCH] BTreeMap: better way to postpone root access in DrainFilter --- library/alloc/src/collections/btree/map.rs | 53 +++++++++------------ library/alloc/src/collections/btree/node.rs | 4 +- 2 files changed, 25 insertions(+), 32 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 1db629c3bdf..75315ceb544 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1697,10 +1697,9 @@ impl<'a, K: 'a, V: 'a> DrainFilterInner<'a, K, V> { let (k, v) = kv.kv_mut(); if pred(k, v) { *self.length -= 1; - let RemoveResult { old_kv, pos, emptied_internal_root } = kv.remove_kv_tracking(); + let (kv, pos) = kv.remove_kv_tracking(|_| self.emptied_internal_root = true); self.cur_leaf_edge = Some(pos); - self.emptied_internal_root |= emptied_internal_root; - return Some(old_kv); + return Some(kv); } self.cur_leaf_edge = Some(kv.next_leaf_edge()); } @@ -2645,35 +2644,28 @@ impl<'a, K: Ord, V> OccupiedEntry<'a, K, V> { fn remove_kv(self) -> (K, V) { *self.length -= 1; - let RemoveResult { old_kv, pos, emptied_internal_root } = self.handle.remove_kv_tracking(); - let root = pos.into_node().into_root_mut(); - if emptied_internal_root { - root.pop_internal_level(); - } + let (old_kv, _) = + self.handle.remove_kv_tracking(|root| root.into_root_mut().pop_internal_level()); old_kv } } -struct RemoveResult<'a, K, V> { - // Key and value removed. - old_kv: (K, V), - // Unique location at the leaf level that the removed KV lopgically collapsed into. - pos: Handle, K, V, marker::Leaf>, marker::Edge>, - // Whether the remove left behind and empty internal root node, that should be removed - // using `pop_internal_level`. - emptied_internal_root: bool, -} - impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInternal>, marker::KV> { /// 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 leaves /// an empty internal root node, which the caller should subsequently pop from /// the map holding the tree. The caller should also decrement the map's length. - fn remove_kv_tracking(self) -> RemoveResult<'a, K, V> { - let (mut pos, old_key, old_val, was_internal) = match self.force() { + fn remove_kv_tracking( + self, + handle_emptied_internal_root: F, + ) -> ((K, V), Handle, K, V, marker::Leaf>, marker::Edge>) + where + F: FnOnce(NodeRef, K, V, marker::Internal>), + { + let (old_kv, mut pos, was_internal) = match self.force() { Leaf(leaf) => { - let (hole, old_key, old_val) = leaf.remove(); - (hole, old_key, old_val, false) + let (old_kv, pos) = leaf.remove(); + (old_kv, pos, false) } Internal(mut internal) => { // Replace the location freed in the internal node with the next KV, @@ -2688,17 +2680,16 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInter let to_remove = internal.left_edge().descend().last_leaf_edge().left_kv().ok(); let to_remove = unsafe { unwrap_unchecked(to_remove) }; - let (hole, key, val) = to_remove.remove(); + let (kv, pos) = to_remove.remove(); - let old_key = unsafe { mem::replace(&mut *key_loc, key) }; - let old_val = unsafe { mem::replace(&mut *val_loc, val) }; + let old_key = unsafe { mem::replace(&mut *key_loc, kv.0) }; + let old_val = unsafe { mem::replace(&mut *val_loc, kv.1) }; - (hole, old_key, old_val, true) + ((old_key, old_val), pos, true) } }; // Handle underflow - let mut emptied_internal_root = false; let mut cur_node = unsafe { ptr::read(&pos).into_node().forget_type() }; let mut at_leaf = true; while cur_node.len() < node::MIN_LEN { @@ -2719,8 +2710,10 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInter let parent = edge.into_node(); if parent.len() == 0 { - // This empty parent must be the root, and should be popped off the tree. - emptied_internal_root = true; + // The parent that was just emptied must be the root, + // because nodes on a lower level would not have been + // left underfull. It has to be popped off the tree soon. + handle_emptied_internal_root(parent); break; } else { cur_node = parent.forget_type(); @@ -2747,7 +2740,7 @@ impl<'a, K: 'a, V: 'a> Handle, K, V, marker::LeafOrInter pos = unsafe { unwrap_unchecked(pos.next_kv().ok()).next_leaf_edge() }; } - RemoveResult { old_kv: (old_key, old_val), pos, emptied_internal_root } + (old_kv, pos) } } diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index 6a4c495ea14..4e52c16d20d 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -1083,12 +1083,12 @@ impl<'a, K, V> Handle, K, V, marker::Leaf>, marker::KV> /// between the now adjacent key/value pairs (if any) to the left and right of this handle. pub fn remove( mut self, - ) -> (Handle, K, V, marker::Leaf>, marker::Edge>, K, V) { + ) -> ((K, V), Handle, K, V, marker::Leaf>, marker::Edge>) { unsafe { let k = slice_remove(self.node.keys_mut(), self.idx); let v = slice_remove(self.node.vals_mut(), self.idx); (*self.node.as_leaf_mut()).len -= 1; - (self.left_edge(), k, v) + ((k, v), self.left_edge()) } } }