Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

BTree: no longer copy keys and values before dropping them #84904

Merged
merged 1 commit into from
May 11, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1439,7 +1439,10 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
impl<K, V> Drop for Dropper<K, V> {
fn drop(&mut self) {
// Similar to advancing a non-fusing iterator.
fn next_or_end<K, V>(this: &mut Dropper<K, V>) -> Option<(K, V)> {
fn next_or_end<K, V>(
this: &mut Dropper<K, V>,
) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>>
{
if this.remaining_length == 0 {
unsafe { ptr::read(&this.front).deallocating_end() }
None
Expand All @@ -1455,13 +1458,15 @@ impl<K, V> Drop for Dropper<K, V> {
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);
}
}
Expand All @@ -1485,7 +1490,9 @@ impl<K, V> Iterator for IntoIter<K, V> {
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())
}
}

Expand All @@ -1501,7 +1508,9 @@ impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
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())
}
}
}
Expand Down
80 changes: 44 additions & 36 deletions library/alloc/src/collections/btree/navigate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -237,25 +237,27 @@ impl<BorrowType: marker::BorrowType, K, V>

impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, 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<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, 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,
Expand All @@ -265,25 +267,27 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, 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<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, 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,
Expand Down Expand Up @@ -373,13 +377,15 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, 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<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
super::mem::replace(self, |leaf_edge| unsafe {
leaf_edge.deallocating_next().unwrap_unchecked()
})
Expand All @@ -391,13 +397,15 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, 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<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
super::mem::replace(self, |leaf_edge| unsafe {
leaf_edge.deallocating_next_back().unwrap_unchecked()
})
Expand Down
39 changes: 36 additions & 3 deletions library/alloc/src/collections/btree/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -422,21 +422,30 @@ impl<'a, K, V, Type> NodeRef<marker::Mut<'a>, 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<K, V> {
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<K, V> {
let ptr = Self::as_leaf_ptr(&mut self);
// SAFETY: we have exclusive access to the entire node.
unsafe { &mut *ptr }
}
}

impl<K, V, Type> NodeRef<marker::Dying, K, V, Type> {
/// Borrows exclusive access to the leaf portion of a dying leaf or internal node.
fn as_leaf_dying(&mut self) -> &mut LeafNode<K, V> {
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<marker::Mut<'a>, K, V, Type> {
/// Borrows exclusive access to an element of the key storage area.
///
Expand Down Expand Up @@ -1040,13 +1049,37 @@ impl<'a, K: 'a, V: 'a, NodeType> Handle<NodeRef<marker::Mut<'a>, 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<K, V, NodeType> Handle<NodeRef<marker::Dying, K, V, NodeType>, 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<NodeRef<marker::Mut<'a>, K, V, NodeType>, marker::KV> {
/// Helps implementations of `split` for a particular `NodeType`,
/// by taking care of leaf data.
Expand Down