Skip to content

BTree: avoid most unsafe code in iterators #86195

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.

Sign up for GitHub

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

Closed
wants to merge 1 commit into from
+118 −135
Closed

BTree: avoid most unsafe code in iterators #86195

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
39 changes: 25 additions & 14 deletions library/alloc/src/collections/btree/map.rs
View file
Original file line number Diff line number Diff line change
@@ -9,7 +9,7 @@ use core::ops::{Index, RangeBounds};
use core::ptr;

use super::borrow::DormantMutRef;
use super::navigate::LeafRange;
use super::navigate::{deallocating_next_unchecked, LeafRange};
use super::node::{self, marker, ForceResult::*, Handle, NodeRef, Root};
use super::search::SearchResult::*;

@@ -149,7 +149,10 @@ pub struct BTreeMap<K, V> {
unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {
fn drop(&mut self) {
if let Some(root) = self.root.take() {
Dropper { front: root.into_dying().first_leaf_edge(), remaining_length: self.length };
Dropper {
front: Some(root.into_dying().first_leaf_edge()),
remaining_length: self.length,
};
}
}
}
@@ -334,7 +337,7 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IntoIter<K, V> {
/// purpose: to drop the remainder of an `IntoIter`. Therefore it also serves to
/// drop an entire tree without the need to first look up a `back` leaf edge.
struct Dropper<K, V> {
front: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
front: Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
remaining_length: usize,
}

@@ -1298,7 +1301,9 @@ impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_unchecked() })
let result = self.range.inner.next_checked();
debug_assert!(result.is_some());
result
}
}

@@ -1329,7 +1334,9 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for Iter<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_back_unchecked() })
let result = self.range.inner.next_back_checked();
debug_assert!(result.is_some());
result
}
}
}
@@ -1367,7 +1374,9 @@ impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_unchecked() })
let result = self.range.inner.next_checked();
debug_assert!(result.is_some());
result
}
}

@@ -1395,7 +1404,9 @@ impl<'a, K: 'a, V: 'a> DoubleEndedIterator for IterMut<'a, K, V> {
None
} else {
self.length -= 1;
Some(unsafe { self.range.inner.next_back_unchecked() })
let result = self.range.inner.next_back_checked();
debug_assert!(result.is_some());
result
}
}
}
@@ -1443,11 +1454,13 @@ impl<K, V> Drop for 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() }
if let Some(front) = this.front.take() {
front.deallocating_end();
}
None
} else {
this.remaining_length -= 1;
Some(unsafe { this.front.deallocating_next_unchecked() })
unsafe { deallocating_next_unchecked(&mut this.front) }
}
}

@@ -1474,9 +1487,7 @@ impl<K, V> Drop for Dropper<K, V> {
#[stable(feature = "btree_drop", since = "1.7.0")]
impl<K, V> Drop for IntoIter<K, V> {
fn drop(&mut self) {
if let Some(front) = self.range.take_front() {
Dropper { front, remaining_length: self.length };
}
Dropper { front: self.range.take_front(), remaining_length: self.length };
}
}

@@ -1490,7 +1501,7 @@ impl<K, V> Iterator for IntoIter<K, V> {
} else {
self.length -= 1;
let kv = unsafe { self.range.deallocating_next_unchecked() };
Some(kv.into_key_val())
Some(kv.unwrap().into_key_val())
}
}

@@ -1507,7 +1518,7 @@ impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
} else {
self.length -= 1;
let kv = unsafe { self.range.deallocating_next_back_unchecked() };
Some(kv.into_key_val())
Some(kv.unwrap().into_key_val())
}
}
}
214 changes: 93 additions & 121 deletions library/alloc/src/collections/btree/navigate.rs
View file
Original file line number Diff line number Diff line change
@@ -35,46 +35,42 @@ impl<BorrowType, K, V> LeafRange<BorrowType, K, V> {
}

impl<'a, K, V> LeafRange<marker::Immut<'a>, K, V> {
/// Moves the leaf edge handle to the next leaf edge and returns
/// references to the key and value in between.
#[inline]
pub fn next_checked(&mut self) -> Option<(&'a K, &'a V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
if self.is_empty() { None } else { perform_next(&mut self.front, |kv| kv.into_kv()) }
}

/// Moves the leaf edge handle to the previous leaf edge and returns
/// references to the key and value in between.
#[inline]
pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
}

#[inline]
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { self.front.as_mut().unwrap().next_unchecked() }
}

#[inline]
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { self.back.as_mut().unwrap().next_back_unchecked() }
if self.is_empty() { None } else { perform_next_back(&mut self.back, |kv| kv.into_kv()) }
}
}

impl<'a, K, V> LeafRange<marker::ValMut<'a>, K, V> {
/// Moves the leaf edge handle to the next leaf edge and returns
/// references to the key and value in between.
#[inline]
pub fn next_checked(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_unchecked() }) }
if self.is_empty() {
None
} else {
perform_next(&mut self.front, |kv| unsafe { ptr::read(kv).into_kv_valmut() })
}
}

/// Moves the leaf edge handle to the previous leaf edge and returns
/// references to the key and value in between.
#[inline]
pub fn next_back_checked(&mut self) -> Option<(&'a K, &'a mut V)> {
if self.is_empty() { None } else { Some(unsafe { self.next_back_unchecked() }) }
}

#[inline]
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
unsafe { self.front.as_mut().unwrap().next_unchecked() }
}

#[inline]
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
unsafe { self.back.as_mut().unwrap().next_back_unchecked() }
if self.is_empty() {
None
} else {
perform_next_back(&mut self.back, |kv| unsafe { ptr::read(kv).into_kv_valmut() })
}
}
}

@@ -89,19 +85,15 @@ impl<K, V> LeafRange<marker::Dying, K, V> {
#[inline]
pub unsafe fn deallocating_next_unchecked(
&mut self,
) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
debug_assert!(self.front.is_some());
let front = self.front.as_mut().unwrap();
unsafe { front.deallocating_next_unchecked() }
) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>> {
unsafe { deallocating_next_unchecked(&mut self.front) }
}

#[inline]
pub unsafe fn deallocating_next_back_unchecked(
&mut self,
) -> Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV> {
debug_assert!(self.back.is_some());
let back = self.back.as_mut().unwrap();
unsafe { back.deallocating_next_back_unchecked() }
) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>> {
unsafe { deallocating_next_back_unchecked(&mut self.back) }
}
}

@@ -388,100 +380,80 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
}
}

impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge> {
/// Moves the leaf edge handle to the next leaf edge and returns references to the
/// key and value in between.
///
/// # Safety
/// There must be another KV in the direction travelled.
unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
super::mem::replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv().ok().unwrap();
(kv.next_leaf_edge(), kv.into_kv())
})
}

/// Moves the leaf edge handle to the previous leaf edge and returns references to the
/// key and value in between.
///
/// # Safety
/// There must be another KV in the direction travelled.
unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
super::mem::replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv().ok().unwrap();
(kv.next_back_leaf_edge(), kv.into_kv())
})
}
/// If possible, extract some result from the next KV and move to the edge beyond it.
fn perform_next<BorrowType: marker::BorrowType, K, V, F, R>(
front: &mut Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
f: F,
) -> Option<R>
where
F: Fn(&Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>) -> R,
{
super::mem::replace(front, |edge| {
if let Some(kv) = edge.and_then(|edge| edge.next_kv().ok()) {
let result = f(&kv);
(Some(kv.next_leaf_edge()), Some(result))
} else {
(None, None)
}
})
}

impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge> {
/// Moves the leaf edge handle to the next leaf edge and returns references to the
/// key and value in between.
///
/// # Safety
/// There must be another KV in the direction travelled.
unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
let kv = super::mem::replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv().ok().unwrap();
(unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
});
// Doing this last is faster, according to benchmarks.
kv.into_kv_valmut()
}

/// Moves the leaf edge handle to the previous leaf and returns references to the
/// key and value in between.
///
/// # Safety
/// There must be another KV in the direction travelled.
unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
let kv = super::mem::replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv().ok().unwrap();
(unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
});
// Doing this last is faster, according to benchmarks.
kv.into_kv_valmut()
}
/// If possible, extract some result from the previous KV and move to the edge beyond it.
fn perform_next_back<BorrowType: marker::BorrowType, K, V, F, R>(
back: &mut Option<Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>>,
f: F,
) -> Option<R>
where
F: Fn(&Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>) -> R,
{
super::mem::replace(back, |edge| {
if let Some(kv) = edge.and_then(|edge| edge.next_back_kv().ok()) {
let result = f(&kv);
(Some(kv.next_back_leaf_edge()), Some(result))
} else {
(None, None)
}
})
}

impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
/// Moves the leaf edge handle to the next leaf edge and returns the key and value
/// in between, deallocating any node left behind while leaving the corresponding
/// edge in its parent node dangling.
///
/// # Safety
/// - There must be another KV in the direction travelled.
/// - 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,
/// 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() })
}
/// Moves the leaf edge handle to the next leaf edge and returns the key and value
/// in between, deallocating any node left behind while leaving the corresponding
/// edge in its parent node dangling.
///
/// # Safety
/// - The next KV, if any, 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,
/// or call this method or counterpart `deallocating_next_back_unchecked` again.
pub unsafe fn deallocating_next_unchecked<K, V>(
opt_hndl: &mut Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>> {
super::mem::replace(opt_hndl, |edge| {
edge.and_then(|edge| unsafe { edge.deallocating_next() })
.map_or((None, None), |p| (Some(p.0), Some(p.1)))
})
}

/// Moves the leaf edge handle to the previous leaf edge and returns the key and value
/// in between, deallocating any node left behind while leaving the corresponding
/// edge in its parent node dangling.
///
/// # Safety
/// - There must be another KV in the direction travelled.
/// - 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,
/// or call this method or counterpart `deallocating_next_unchecked` again.
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()
})
}
/// Moves the leaf edge handle to the previous leaf edge and returns the key and value
/// in between, deallocating any node left behind while leaving the corresponding
/// edge in its parent node dangling.
///
/// # Safety
/// - The previous KV, if any, 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,
/// or call this method or counterpart `deallocating_next_unchecked` again.
unsafe fn deallocating_next_back_unchecked<K, V>(
opt_hndl: &mut Option<Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>>,
) -> Option<Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>> {
super::mem::replace(opt_hndl, |edge| {
edge.and_then(|edge| unsafe { edge.deallocating_next_back() })
.map_or((None, None), |p| (Some(p.0), Some(p.1)))
})
}

impl<BorrowType: marker::BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {