diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs index bbeced1751d..aedba799169 100644 --- a/src/liballoc/collections/btree/map.rs +++ b/src/liballoc/collections/btree/map.rs @@ -1780,18 +1780,12 @@ pub(super) fn next(&mut self, pred: &mut F) -> Option<(K, V)> where F: FnMut(&K, &mut V) -> bool, { - while let Some(kv) = unsafe { self.next_kv() } { - let (k, v) = unsafe { ptr::read(&kv) }.into_kv_mut(); + while let Some(mut kv) = unsafe { self.next_kv() } { + let (k, v) = kv.kv_mut(); if pred(k, v) { *self.length -= 1; let (k, v, leaf_edge_location) = kv.remove_kv_tracking(); - // `remove_kv_tracking` has either preserved or invalidated `self.cur_leaf_edge` - if let Some(node) = leaf_edge_location { - match search::search_tree(node, &k) { - search::SearchResult::Found(_) => unreachable!(), - search::SearchResult::GoDown(leaf) => self.cur_leaf_edge = Some(leaf), - } - }; + self.cur_leaf_edge = Some(leaf_edge_location); return Some((k, v)); } self.cur_leaf_edge = Some(kv.next_leaf_edge()); @@ -2698,28 +2692,26 @@ fn remove_kv(self) -> (K, V) { 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 whereabouts of the leaf edge corresponding to that former pair: - /// if None is returned, the leaf edge is still the left leaf edge of the KV handle; - /// if a node is returned, it heads the subtree where the leaf edge may be found. + /// the leaf edge corresponding to that former pair. fn remove_kv_tracking( self, - ) -> (K, V, Option, K, V, marker::LeafOrInternal>>) { - let mut levels_down_handled: isize; - let (small_leaf, old_key, old_val) = match self.force() { + ) -> (K, V, Handle, K, V, marker::Leaf>, marker::Edge>) { + let (mut pos, old_key, old_val, was_internal) = match self.force() { Leaf(leaf) => { - levels_down_handled = 1; // handled at same level, but affects only the right side let (hole, old_key, old_val) = leaf.remove(); - (hole.into_node(), old_key, old_val) + (hole, old_key, old_val, false) } Internal(mut internal) => { // Replace the location freed in the internal node with the next KV, // and remove that next KV from its leaf. - levels_down_handled = unsafe { ptr::read(&internal).into_node().height() } as isize; let key_loc = internal.kv_mut().0 as *mut K; let val_loc = internal.kv_mut().1 as *mut V; - let to_remove = internal.right_edge().descend().first_leaf_edge().right_kv().ok(); + // Deleting from the left side is typically faster since we can + // just pop an element from the end of the KV array without + // needing to shift the other values. + 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(); @@ -2727,50 +2719,72 @@ fn remove_kv_tracking( let old_key = unsafe { mem::replace(&mut *key_loc, key) }; let old_val = unsafe { mem::replace(&mut *val_loc, val) }; - (hole.into_node(), old_key, old_val) + (hole, old_key, old_val, true) } }; // Handle underflow - let mut cur_node = small_leaf.forget_type(); + 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(root) => { - cur_node = root; - break; - } + AtRoot(_) => break, EmptyParent(_) => unreachable!(), - Merged(parent) => { - levels_down_handled -= 1; + 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!(), + }; + debug_assert!(idx <= node.len()); + pos = unsafe { Handle::new_edge(node, idx) }; + } + + let parent = edge.into_node(); if parent.len() == 0 { // We must be at the root - let root = parent.into_root_mut(); - root.pop_level(); - cur_node = root.as_mut(); + parent.into_root_mut().pop_level(); break; } else { cur_node = parent.forget_type(); + at_leaf = false; } } - Stole(internal_node) => { - levels_down_handled -= 1; - cur_node = internal_node.forget_type(); + 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.next_unchecked(); + } + } + // This internal node might be underfull, but only if it's the root. break; } } } - let leaf_edge_location = if levels_down_handled > 0 { None } else { Some(cur_node) }; - (old_key, old_val, leaf_edge_location) + // 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_key, old_val, pos) } } enum UnderflowResult<'a, K, V> { AtRoot(NodeRef, K, V, marker::LeafOrInternal>), EmptyParent(NodeRef, K, V, marker::Internal>), - Merged(NodeRef, K, V, marker::Internal>), - Stole(NodeRef, K, V, marker::Internal>), + Merged(Handle, K, V, marker::Internal>, marker::Edge>, bool, usize), + Stole(NodeRef, K, V, marker::Internal>, bool), } fn handle_underfull_node( @@ -2792,14 +2806,15 @@ fn handle_underfull_node( }; if handle.can_merge() { - Merged(handle.merge().into_node()) + 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(handle.into_node()) + Stole(handle.into_node(), is_left) } } diff --git a/src/liballoc/collections/btree/node.rs b/src/liballoc/collections/btree/node.rs index 11c14299573..bc4e2711670 100644 --- a/src/liballoc/collections/btree/node.rs +++ b/src/liballoc/collections/btree/node.rs @@ -723,6 +723,11 @@ impl Handle { pub fn into_node(self) -> Node { self.node } + + /// Returns the position of this handle in the node. + pub fn idx(&self) -> usize { + self.idx + } } impl Handle, marker::KV> {