Skip to content

Commit

Permalink
Auto merge of #78359 - ssomers:btree_cleanup_mem, r=Mark-Simulacrum
Browse files Browse the repository at this point in the history
BTreeMap: move generic support functions out of navigate.rs

A preparatory step chipped off #78104, useful in general (if at all).

r? `@Mark-Simulacrum`
  • Loading branch information
bors committed Oct 27, 2020
2 parents 28f03ac + 0da7941 commit c9b606e
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 40 deletions.
34 changes: 34 additions & 0 deletions library/alloc/src/collections/btree/mem.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
use core::intrinsics;
use core::mem;
use core::ptr;

/// This replaces the value behind the `v` unique reference by calling the
/// relevant function.
///
/// If a panic occurs in the `change` closure, the entire process will be aborted.
#[inline]
pub fn take_mut<T>(v: &mut T, change: impl FnOnce(T) -> T) {
replace(v, |value| (change(value), ()))
}

/// This replaces the value behind the `v` unique reference by calling the
/// relevant function, and returns a result obtained along the way.
///
/// If a panic occurs in the `change` closure, the entire process will be aborted.
#[inline]
pub 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
}
1 change: 1 addition & 0 deletions library/alloc/src/collections/btree/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
mod borrow;
pub mod map;
mod mem;
mod merge_iter;
mod navigate;
mod node;
Expand Down
47 changes: 7 additions & 40 deletions library/alloc/src/collections/btree/navigate.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
use core::borrow::Borrow;
use core::cmp::Ordering;
use core::intrinsics;
use core::mem;
use core::ops::Bound::{Excluded, Included, Unbounded};
use core::ops::RangeBounds;
use core::ptr;
Expand Down Expand Up @@ -304,45 +302,14 @@ macro_rules! def_next_kv_uncheched_dealloc {
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.
///
/// If a panic occurs in the `change` closure, the entire process will be aborted.
#[inline]
fn take_mut<T>(v: &mut T, change: impl FnOnce(T) -> T) {
replace(v, |value| (change(value), ()))
}

/// This replaces the value behind the `v` unique reference by calling the
/// relevant function, and returns a result obtained along the way.
///
/// If a panic occurs in the `change` closure, the entire process will be aborted.
#[inline]
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
}

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.
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
replace(self, |leaf_edge| {
super::mem::replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(kv.next_leaf_edge(), kv.into_kv())
Expand All @@ -355,7 +322,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Ed
/// # Safety
/// There must be another KV in the direction travelled.
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
replace(self, |leaf_edge| {
super::mem::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 @@ -370,7 +337,7 @@ impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::E
/// # Safety
/// There must be another KV in the direction travelled.
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
let kv = replace(self, |leaf_edge| {
let kv = super::mem::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)
Expand All @@ -385,7 +352,7 @@ impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::E
/// # Safety
/// There must be another KV in the direction travelled.
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
let kv = replace(self, |leaf_edge| {
let kv = super::mem::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)
Expand All @@ -401,7 +368,7 @@ impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge
/// # Safety
/// There must be another KV in the direction travelled.
pub unsafe fn move_next_unchecked(&mut self) {
take_mut(self, |leaf_edge| {
super::mem::take_mut(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
kv.next_leaf_edge()
Expand All @@ -423,7 +390,7 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
/// call this method again subject to its safety conditions, or call counterpart
/// `next_back_unchecked` subject to its safety conditions.
pub unsafe fn next_unchecked(&mut self) -> (K, V) {
replace(self, |leaf_edge| {
super::mem::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) };
Expand All @@ -444,7 +411,7 @@ impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
/// call this method again subject to its safety conditions, or call counterpart
/// `next_unchecked` subject to its safety conditions.
pub unsafe fn next_back_unchecked(&mut self) -> (K, V) {
replace(self, |leaf_edge| {
super::mem::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) };
Expand Down

0 comments on commit c9b606e

Please sign in to comment.