squiid/rust
1
0
Fork 0
Code Issues Pull requests Projects Activity

Auto merge of #75071 - ssomers:btree_cleanup_5, r=Mark-Simulacrum

Browse source 
BTreeMap: enforce the panic rule imposed by `replace`

Also, reveal the unsafe parts in the closures fed to it.

r? @Mark-Simulacrum
This commit is contained in:
bors 2020-08-07 21:48:32 +00:00
parent 09f4c9f508 734fc0477c
commit c2d1b0d980
1 changed files with 51 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

104
library/alloc/src/collections/btree/navigate.rs
View file

@ -1,3 +1,5 @@
use core::intrinsics;
use core::mem;
use core::ptr; use core::ptr;
use super::node::{marker, ForceResult::*, Handle, NodeRef}; use super::node::{marker, ForceResult::*, Handle, NodeRef};
@ -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} 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 /// 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] #[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 value = unsafe { ptr::read(v) };
let (new_value, ret) = change(value); let (new_value, ret) = change(value);
unsafe { unsafe {
ptr::write(v, new_value); ptr::write(v, new_value);
} }
mem::forget(guard);
ret ret
} }
@ -97,26 +107,22 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
/// key and value in between. /// key and value in between.
/// Unsafe because the caller must ensure that the leaf edge is not the last one in the tree. /// 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) { pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { replace(self, |leaf_edge| {
replace(self, |leaf_edge| { let kv = leaf_edge.next_kv();
let kv = leaf_edge.next_kv(); let kv = unsafe { unwrap_unchecked(kv.ok()) };
let kv = unwrap_unchecked(kv.ok()); (kv.next_leaf_edge(), kv.into_kv())
(kv.next_leaf_edge(), kv.into_kv()) })
})
}
} }
/// Moves the leaf edge handle to the previous leaf edge and returns references to the /// Moves the leaf edge handle to the previous leaf edge and returns references to the
/// key and value in between. /// key and value in between.
/// Unsafe because the caller must ensure that the leaf edge is not the first one in the tree. /// 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) { pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
unsafe { replace(self, |leaf_edge| {
replace(self, |leaf_edge| { let kv = leaf_edge.next_back_kv();
let kv = leaf_edge.next_back_kv(); let kv = unsafe { unwrap_unchecked(kv.ok()) };
let kv = unwrap_unchecked(kv.ok()); (kv.next_back_leaf_edge(), kv.into_kv())
(kv.next_back_leaf_edge(), kv.into_kv()) })
})
}
} }
} }
@ -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. /// - 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. /// - Using the updated handle may well invalidate the returned references.
pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) { pub unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
unsafe { let kv = replace(self, |leaf_edge| {
let kv = replace(self, |leaf_edge| { let kv = leaf_edge.next_kv();
let kv = leaf_edge.next_kv(); let kv = unsafe { unwrap_unchecked(kv.ok()) };
let kv = unwrap_unchecked(kv.ok()); (unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
(ptr::read(&kv).next_leaf_edge(), kv) });
}); // Doing the descend (and perhaps another move) invalidates the references
// Doing the descend (and perhaps another move) invalidates the references // returned by `into_kv_mut`, so we have to do this last.
// returned by `into_kv_mut`, so we have to do this last. kv.into_kv_mut()
kv.into_kv_mut()
}
} }
/// Moves the leaf edge handle to the previous leaf and returns references to the /// Moves the leaf edge handle to the previous leaf and returns references to the
@ -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. /// - 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. /// - Using the updated handle may well invalidate the returned references.
pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) { pub unsafe fn next_back_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
unsafe { let kv = replace(self, |leaf_edge| {
let kv = replace(self, |leaf_edge| { let kv = leaf_edge.next_back_kv();
let kv = leaf_edge.next_back_kv(); let kv = unsafe { unwrap_unchecked(kv.ok()) };
let kv = unwrap_unchecked(kv.ok()); (unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
(ptr::read(&kv).next_back_leaf_edge(), kv) });
}); // Doing the descend (and perhaps another move) invalidates the references
// Doing the descend (and perhaps another move) invalidates the references // returned by `into_kv_mut`, so we have to do this last.
// returned by `into_kv_mut`, so we have to do this last. kv.into_kv_mut()
kv.into_kv_mut()
}
} }
} }
@ -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, /// call this method again subject to both preconditions listed in the first point,
/// or call counterpart `next_back_unchecked` subject to its preconditions. /// or call counterpart `next_back_unchecked` subject to its preconditions.
pub unsafe fn next_unchecked(&mut self) -> (K, V) { pub unsafe fn next_unchecked(&mut self) -> (K, V) {
unsafe { replace(self, |leaf_edge| {
replace(self, |leaf_edge| { let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge) };
let kv = next_kv_unchecked_dealloc(leaf_edge); let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
let k = ptr::read(kv.reborrow().into_kv().0); let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
let v = ptr::read(kv.reborrow().into_kv().1); (kv.next_leaf_edge(), (k, v))
(kv.next_leaf_edge(), (k, v)) })
})
}
} }
/// Moves the leaf edge handle to the previous leaf edge and returns the key /// Moves the leaf edge handle to the previous leaf edge and returns the key
@ -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, /// call this method again subject to both preconditions listed in the first point,
/// or call counterpart `next_unchecked` subject to its preconditions. /// or call counterpart `next_unchecked` subject to its preconditions.
pub unsafe fn next_back_unchecked(&mut self) -> (K, V) { pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
unsafe { replace(self, |leaf_edge| {
replace(self, |leaf_edge| { let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge) };
let kv = next_back_kv_unchecked_dealloc(leaf_edge); let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
let k = ptr::read(kv.reborrow().into_kv().0); let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
let v = ptr::read(kv.reborrow().into_kv().1); (kv.next_back_leaf_edge(), (k, v))
(kv.next_back_leaf_edge(), (k, v)) })
})
}
} }
} }