From 728204b40efe11ecc866387924a57c684786e2b9 Mon Sep 17 00:00:00 2001 From: Stein Somers Date: Mon, 3 May 2021 11:54:15 +0200 Subject: [PATCH] BTree: no longer copy keys and values before dropping them --- library/alloc/src/collections/btree/map.rs | 21 +++-- .../alloc/src/collections/btree/navigate.rs | 80 ++++++++++--------- library/alloc/src/collections/btree/node.rs | 39 ++++++++- 3 files changed, 95 insertions(+), 45 deletions(-) diff --git a/library/alloc/src/collections/btree/map.rs b/library/alloc/src/collections/btree/map.rs index 0a46387c34e9c..30194aa446f4b 100644 --- a/library/alloc/src/collections/btree/map.rs +++ b/library/alloc/src/collections/btree/map.rs @@ -1439,7 +1439,10 @@ impl IntoIterator for BTreeMap { impl Drop for Dropper { fn drop(&mut self) { // Similar to advancing a non-fusing iterator. - fn next_or_end(this: &mut Dropper) -> Option<(K, V)> { + fn next_or_end( + this: &mut Dropper, + ) -> Option, marker::KV>> + { if this.remaining_length == 0 { unsafe { ptr::read(&this.front).deallocating_end() } None @@ -1455,13 +1458,15 @@ impl Drop for Dropper { fn drop(&mut self) { // Continue the same loop we perform below. This only runs when unwinding, so we // don't have to care about panics this time (they'll abort). - while let Some(_pair) = next_or_end(&mut self.0) {} + while let Some(kv) = next_or_end(&mut self.0) { + kv.drop_key_val(); + } } } - while let Some(pair) = next_or_end(self) { + while let Some(kv) = next_or_end(self) { let guard = DropGuard(self); - drop(pair); + kv.drop_key_val(); mem::forget(guard); } } @@ -1485,7 +1490,9 @@ impl Iterator for IntoIter { None } else { self.length -= 1; - Some(unsafe { self.range.front.as_mut().unwrap().deallocating_next_unchecked() }) + let front = self.range.front.as_mut().unwrap(); + let kv = unsafe { front.deallocating_next_unchecked() }; + Some(kv.into_key_val()) } } @@ -1501,7 +1508,9 @@ impl DoubleEndedIterator for IntoIter { None } else { self.length -= 1; - Some(unsafe { self.range.back.as_mut().unwrap().deallocating_next_back_unchecked() }) + let back = self.range.back.as_mut().unwrap(); + let kv = unsafe { back.deallocating_next_back_unchecked() }; + Some(kv.into_key_val()) } } } diff --git a/library/alloc/src/collections/btree/navigate.rs b/library/alloc/src/collections/btree/navigate.rs index 4399feaccc9b7..563c070dd0f1f 100644 --- a/library/alloc/src/collections/btree/navigate.rs +++ b/library/alloc/src/collections/btree/navigate.rs @@ -237,25 +237,27 @@ impl impl Handle, marker::Edge> { /// Given a leaf edge handle into a dying tree, returns the next leaf edge - /// on the right side, and the key-value pair in between, which is either - /// in the same leaf node, in an ancestor node, or non-existent. + /// on the right side, and the key-value pair in between, if they exist. /// - /// This method also deallocates any node(s) it reaches the end of. This - /// implies that if no more key-value pair exists, the entire remainder of - /// the tree will have been deallocated and there is nothing left to return. + /// If the given edge is the last one in a leaf, this method deallocates + /// the leaf, as well as any ancestor nodes whose last edge was reached. + /// This implies that if no more key-value pair follows, the entire tree + /// will have been deallocated and there is nothing left to return. /// /// # Safety - /// The given edge must not have been previously returned by counterpart - /// `deallocating_next_back`. - unsafe fn deallocating_next(self) -> Option<(Self, (K, V))> { + /// - The given edge must not have been previously returned by counterpart + /// `deallocating_next_back`. + /// - The returned KV handle is only valid to access the key and value, + /// and only valid until the next call to this method or counterpart + /// `deallocating_next_back`. + pub unsafe fn deallocating_next( + self, + ) -> Option<(Self, Handle, marker::KV>)> + { let mut edge = self.forget_node_type(); loop { edge = match edge.right_kv() { - Ok(kv) => { - let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; - let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; - return Some((kv.next_leaf_edge(), (k, v))); - } + Ok(kv) => return Some((unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)), Err(last_edge) => match unsafe { last_edge.into_node().deallocate_and_ascend() } { Some(parent_edge) => parent_edge.forget_node_type(), None => return None, @@ -265,25 +267,27 @@ impl Handle, marker::Edge> { } /// Given a leaf edge handle into a dying tree, returns the next leaf edge - /// on the left side, and the key-value pair in between, which is either - /// in the same leaf node, in an ancestor node, or non-existent. + /// on the left side, and the key-value pair in between, if they exist. /// - /// This method also deallocates any node(s) it reaches the end of. This - /// implies that if no more key-value pair exists, the entire remainder of - /// the tree will have been deallocated and there is nothing left to return. + /// If the given edge is the first one in a leaf, this method deallocates + /// the leaf, as well as any ancestor nodes whose first edge was reached. + /// This implies that if no more key-value pair follows, the entire tree + /// will have been deallocated and there is nothing left to return. /// /// # Safety - /// The given edge must not have been previously returned by counterpart - /// `deallocating_next`. - unsafe fn deallocating_next_back(self) -> Option<(Self, (K, V))> { + /// - The given edge must not have been previously returned by counterpart + /// `deallocating_next`. + /// - The returned KV handle is only valid to access the key and value, + /// and only valid until the next call to this method or counterpart + /// `deallocating_next`. + unsafe fn deallocating_next_back( + self, + ) -> Option<(Self, Handle, marker::KV>)> + { let mut edge = self.forget_node_type(); loop { edge = match edge.left_kv() { - Ok(kv) => { - let k = unsafe { ptr::read(kv.reborrow().into_kv().0) }; - let v = unsafe { ptr::read(kv.reborrow().into_kv().1) }; - return Some((kv.next_back_leaf_edge(), (k, v))); - } + Ok(kv) => return Some((unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)), Err(last_edge) => match unsafe { last_edge.into_node().deallocate_and_ascend() } { Some(parent_edge) => parent_edge.forget_node_type(), None => return None, @@ -373,13 +377,15 @@ impl Handle, marker::Edge> { /// /// # Safety /// - There must be another KV in the direction travelled. - /// - That KV was not previously returned by counterpart `next_back_unchecked` - /// on any copy of the handles being used to traverse the tree. + /// - That KV was not previously returned by counterpart + /// `deallocating_next_back_unchecked` on any copy of the handles + /// being used to traverse the tree. /// /// The only safe way to proceed with the updated handle is to compare it, drop it, - /// call this method again subject to its safety conditions, or call counterpart - /// `next_back_unchecked` subject to its safety conditions. - pub unsafe fn deallocating_next_unchecked(&mut self) -> (K, V) { + /// or call this method or counterpart `deallocating_next_back_unchecked` again. + pub unsafe fn deallocating_next_unchecked( + &mut self, + ) -> Handle, marker::KV> { super::mem::replace(self, |leaf_edge| unsafe { leaf_edge.deallocating_next().unwrap_unchecked() }) @@ -391,13 +397,15 @@ impl Handle, marker::Edge> { /// /// # Safety /// - There must be another KV in the direction travelled. - /// - That leaf edge was not previously returned by counterpart `next_unchecked` - /// on any copy of the handles being used to traverse the tree. + /// - That leaf edge was not previously returned by counterpart + /// `deallocating_next_unchecked` on any copy of the handles + /// being used to traverse the tree. /// /// The only safe way to proceed with the updated handle is to compare it, drop it, - /// call this method again subject to its safety conditions, or call counterpart - /// `next_unchecked` subject to its safety conditions. - pub unsafe fn deallocating_next_back_unchecked(&mut self) -> (K, V) { + /// or call this method or counterpart `deallocating_next_unchecked` again. + pub unsafe fn deallocating_next_back_unchecked( + &mut self, + ) -> Handle, marker::KV> { super::mem::replace(self, |leaf_edge| unsafe { leaf_edge.deallocating_next_back().unwrap_unchecked() }) diff --git a/library/alloc/src/collections/btree/node.rs b/library/alloc/src/collections/btree/node.rs index af403496e38e9..3c453529ba84d 100644 --- a/library/alloc/src/collections/btree/node.rs +++ b/library/alloc/src/collections/btree/node.rs @@ -422,14 +422,14 @@ impl<'a, K, V, Type> NodeRef, K, V, Type> { NodeRef { height: self.height, node: self.node, _marker: PhantomData } } - /// Borrows exclusive access to the leaf portion of any leaf or internal node. + /// Borrows exclusive access to the leaf portion of a leaf or internal node. fn as_leaf_mut(&mut self) -> &mut LeafNode { let ptr = Self::as_leaf_ptr(self); // SAFETY: we have exclusive access to the entire node. unsafe { &mut *ptr } } - /// Offers exclusive access to the leaf portion of any leaf or internal node. + /// Offers exclusive access to the leaf portion of a leaf or internal node. fn into_leaf_mut(mut self) -> &'a mut LeafNode { let ptr = Self::as_leaf_ptr(&mut self); // SAFETY: we have exclusive access to the entire node. @@ -437,6 +437,15 @@ impl<'a, K, V, Type> NodeRef, K, V, Type> { } } +impl NodeRef { + /// Borrows exclusive access to the leaf portion of a dying leaf or internal node. + fn as_leaf_dying(&mut self) -> &mut LeafNode { + let ptr = Self::as_leaf_ptr(self); + // SAFETY: we have exclusive access to the entire node. + unsafe { &mut *ptr } + } +} + impl<'a, K: 'a, V: 'a, Type> NodeRef, K, V, Type> { /// Borrows exclusive access to an element of the key storage area. /// @@ -1040,13 +1049,37 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType> } } - /// Replace the key and value that the KV handle refers to. + /// Replaces the key and value that the KV handle refers to. pub fn replace_kv(&mut self, k: K, v: V) -> (K, V) { let (key, val) = self.kv_mut(); (mem::replace(key, k), mem::replace(val, v)) } } +impl Handle, marker::KV> { + /// Extracts the key and value that the KV handle refers to. + pub fn into_key_val(mut self) -> (K, V) { + debug_assert!(self.idx < self.node.len()); + let leaf = self.node.as_leaf_dying(); + unsafe { + let key = leaf.keys.get_unchecked_mut(self.idx).assume_init_read(); + let val = leaf.vals.get_unchecked_mut(self.idx).assume_init_read(); + (key, val) + } + } + + /// Drops the key and value that the KV handle refers to. + #[inline] + pub fn drop_key_val(mut self) { + debug_assert!(self.idx < self.node.len()); + let leaf = self.node.as_leaf_dying(); + unsafe { + leaf.keys.get_unchecked_mut(self.idx).assume_init_drop(); + leaf.vals.get_unchecked_mut(self.idx).assume_init_drop(); + } + } +} + impl<'a, K: 'a, V: 'a, NodeType> Handle, K, V, NodeType>, marker::KV> { /// Helps implementations of `split` for a particular `NodeType`, /// by taking care of leaf data.