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

BTreeMap: enforce the panic rule imposed by replace #75071

Merged
merged 1 commit into from
Aug 8, 2020
Merged
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
104 changes: 51 additions & 53 deletions library/alloc/src/collections/btree/navigate.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use core::intrinsics;
use core::mem;
use core::ptr;

use super::node::{marker, ForceResult::*, Handle, NodeRef};
Expand Down Expand Up @@ -79,16 +81,24 @@ def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}

/// This replaces the value behind the `v` unique reference by calling the
/// relevant function.
/// relevant function, and returns a result obtained along the way.
///
/// Safety: The change closure must not panic.
/// If a panic occurs in the `change` closure, the entire process will be aborted.
#[inline]
unsafe fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
struct PanicGuard;
impl Drop for PanicGuard {
fn drop(&mut self) {
intrinsics::abort()
}
}
let guard = PanicGuard;
let value = unsafe { ptr::read(v) };
let (new_value, ret) = change(value);
unsafe {
ptr::write(v, new_value);
}
mem::forget(guard);
ret
}

Expand All @@ -97,26 +107,22 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
/// key and value in between.
/// Unsafe because the caller must ensure that the leaf edge is not the last one in the tree.
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe {
replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unwrap_unchecked(kv.ok());
(kv.next_leaf_edge(), kv.into_kv())
})
}
replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(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.
/// Unsafe because the caller must ensure that the leaf edge is not the first one in the tree.
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe {
replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv();
let kv = unwrap_unchecked(kv.ok());
(kv.next_back_leaf_edge(), kv.into_kv())
})
}
replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(kv.next_back_leaf_edge(), kv.into_kv())
})
}
}

Expand All @@ -127,16 +133,14 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
/// - The caller must ensure that the leaf edge is not the last one in the tree.
/// - Using the updated handle may well invalidate the returned references.
pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
unsafe {
let kv = replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unwrap_unchecked(kv.ok());
(ptr::read(&kv).next_leaf_edge(), kv)
});
// Doing the descend (and perhaps another move) invalidates the references
// returned by `into_kv_mut`, so we have to do this last.
kv.into_kv_mut()
}
let kv = replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
});
// Doing the descend (and perhaps another move) invalidates the references
// returned by `into_kv_mut`, so we have to do this last.
kv.into_kv_mut()
}

/// Moves the leaf edge handle to the previous leaf and returns references to the
Expand All @@ -145,16 +149,14 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
/// - The caller must ensure that the leaf edge is not the first one in the tree.
/// - Using the updated handle may well invalidate the returned references.
pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
unsafe {
let kv = replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv();
let kv = unwrap_unchecked(kv.ok());
(ptr::read(&kv).next_back_leaf_edge(), kv)
});
// Doing the descend (and perhaps another move) invalidates the references
// returned by `into_kv_mut`, so we have to do this last.
kv.into_kv_mut()
}
let kv = replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
});
// Doing the descend (and perhaps another move) invalidates the references
// returned by `into_kv_mut`, so we have to do this last.
kv.into_kv_mut()
}
}

Expand All @@ -172,14 +174,12 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
/// call this method again subject to both preconditions listed in the first point,
/// or call counterpart `next_back_unchecked` subject to its preconditions.
pub unsafe fn next_unchecked(&mut self) -> (K, V) {
unsafe {
replace(self, |leaf_edge| {
let kv = next_kv_unchecked_dealloc(leaf_edge);
let k = ptr::read(kv.reborrow().into_kv().0);
let v = ptr::read(kv.reborrow().into_kv().1);
(kv.next_leaf_edge(), (k, v))
})
}
replace(self, |leaf_edge| {
let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge) };
let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
(kv.next_leaf_edge(), (k, v))
})
}

/// Moves the leaf edge handle to the previous leaf edge and returns the key
Expand All @@ -195,14 +195,12 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
/// call this method again subject to both preconditions listed in the first point,
/// or call counterpart `next_unchecked` subject to its preconditions.
pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
unsafe {
replace(self, |leaf_edge| {
let kv = next_back_kv_unchecked_dealloc(leaf_edge);
let k = ptr::read(kv.reborrow().into_kv().0);
let v = ptr::read(kv.reborrow().into_kv().1);
(kv.next_back_leaf_edge(), (k, v))
})
}
replace(self, |leaf_edge| {
let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge) };
let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
(kv.next_back_leaf_edge(), (k, v))
})
}
}

Expand Down