Skip to content
Closed
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
17 changes: 14 additions & 3 deletions library/alloc/src/collections/btree/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
// an edge both identifies a position and contains a pointer to a child node.

use core::marker::PhantomData;
use core::mem::{self, DropGuard, MaybeUninit};
use core::mem::{self, MaybeUninit};
use core::num::NonZero;
use core::ptr::{self, NonNull};
use core::slice::SliceIndex;
Expand Down Expand Up @@ -1237,14 +1237,25 @@ impl<K, V, NodeType> Handle<NodeRef<marker::Dying, K, V, NodeType>, marker::KV>
/// The node that the handle refers to must not yet have been deallocated.
#[inline]
pub(super) unsafe fn drop_key_val(mut self) {
// Run the destructor of the value even if the destructor of the key panics.
struct Dropper<'a, T>(&'a mut MaybeUninit<T>);
impl<T> Drop for Dropper<'_, T> {
#[inline]
fn drop(&mut self) {
// ignore-tidy-undocumented-unsafe
unsafe {
self.0.assume_init_drop();
}
}
}

debug_assert!(self.idx < self.node.len());
let leaf = self.node.as_leaf_dying();
// ignore-tidy-undocumented-unsafe
unsafe {
let key = leaf.keys.get_unchecked_mut(self.idx);
let val = leaf.vals.get_unchecked_mut(self.idx);
// Run the destructor of the value even if the destructor of the key panics.
let _guard = DropGuard::new(val, |val| val.assume_init_drop());
let _guard = Dropper(val);
key.assume_init_drop();
// dropping the guard will drop the value
}
Expand Down
247 changes: 127 additions & 120 deletions library/alloc/src/collections/vec_deque/drain.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use core::iter::FusedIterator;
use core::marker::PhantomData;
use core::mem::{self, DropGuard, SizedTypeProperties};
use core::mem::{self, SizedTypeProperties};
use core::ptr::NonNull;
use core::{fmt, ptr};

Expand Down Expand Up @@ -94,137 +94,144 @@ unsafe impl<T: Send, A: Allocator + Send> Send for Drain<'_, T, A> {}
#[stable(feature = "drain", since = "1.6.0")]
impl<T, A: Allocator> Drop for Drain<'_, T, A> {
fn drop(&mut self) {
struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>);

let guard = DropGuard(self);

if mem::needs_drop::<T>() && guard.0.remaining != 0 {
// SAFETY: We just checked that `self.remaining != 0`.
let (front, back) = unsafe { guard.0.as_slices() };
// since idx is a logical index, we don't need to worry about wrapping.
guard.0.idx += front.len();
guard.0.remaining -= front.len();
// SAFETY: This can't have been dropped before since
// `idx` & `remaining` track what's been dropped.
unsafe { ptr::drop_in_place(front) };
guard.0.remaining = 0;
// SAFETY: Ditto.
unsafe { ptr::drop_in_place(back) };
}

// Dropping `guard` handles moving the remaining elements into place.
let mut guard = DropGuard::new(self, |drain| {
if mem::needs_drop::<T>() && drain.remaining != 0 {
// SAFETY: We just checked that `self.remaining != 0`.
unsafe {
let (front, back) = drain.as_slices();
ptr::drop_in_place(front);
ptr::drop_in_place(back);
impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
#[inline]
fn drop(&mut self) {
if mem::needs_drop::<T>() && self.0.remaining != 0 {
// SAFETY: We just checked that `self.remaining != 0`.
unsafe {
let (front, back) = self.0.as_slices();
ptr::drop_in_place(front);
ptr::drop_in_place(back);
}
}
}

// ignore-tidy-undocumented-unsafe
let source_deque = unsafe { drain.deque.as_mut() };
// ignore-tidy-undocumented-unsafe
let source_deque = unsafe { self.0.deque.as_mut() };

let drain_len = drain.drain_len;
let head_len = source_deque.len; // #elements in front of the drain
let tail_len = drain.tail_len; // #elements behind the drain
let new_len = head_len + tail_len;
let drain_len = self.0.drain_len;
let head_len = source_deque.len; // #elements in front of the drain
let tail_len = self.0.tail_len; // #elements behind the drain
let new_len = head_len + tail_len;

if T::IS_ZST {
// no need to copy around any memory if T is a ZST
source_deque.len = new_len;
return;
}
if T::IS_ZST {
// no need to copy around any memory if T is a ZST
source_deque.len = new_len;
return;
}

// Next, we will fill the hole left by the drain with as few writes as possible.
// The code below handles the following control flow and reduces the amount of
// branches under the assumption that `head_len == 0 || tail_len == 0`, i.e.
// draining at the front or at the back of the dequeue is especially common.
//
// H = "head index" = `deque.head`
// h = elements in front of the drain
// d = elements in the drain
// t = elements behind the drain
//
// Note that the buffer may wrap at any point and the wrapping is handled by
// `wrap_copy` and `to_physical_idx`.
//
// Case 1: if `head_len == 0 && tail_len == 0`
// Everything was drained, reset the head index back to 0.
// H
// [ . . . . . d d d d . . . . . ]
// H
// [ . . . . . . . . . . . . . . ]
//
// Case 2: else if `tail_len == 0`
// Don't move data or the head index.
// H
// [ . . . h h h h d d d d . . . ]
// H
// [ . . . h h h h . . . . . . . ]
//
// Case 3: else if `head_len == 0`
// Don't move data, but move the head index.
// H
// [ . . . d d d d t t t t . . . ]
// H
// [ . . . . . . . t t t t . . . ]
//
// Case 4: else if `tail_len <= head_len`
// Move data, but not the head index.
// H
// [ . . h h h h d d d d t t . . ]
// H
// [ . . h h h h t t . . . . . . ]
//
// Case 5: else
// Move data and the head index.
// H
// [ . . h h d d d d t t t t . . ]
// H
// [ . . . . . . h h t t t t . . ]
// Next, we will fill the hole left by the drain with as few writes as possible.
// The code below handles the following control flow and reduces the amount of
// branches under the assumption that `head_len == 0 || tail_len == 0`, i.e.
// draining at the front or at the back of the dequeue is especially common.
//
// H = "head index" = `deque.head`
// h = elements in front of the drain
// d = elements in the drain
// t = elements behind the drain
//
// Note that the buffer may wrap at any point and the wrapping is handled by
// `wrap_copy` and `to_physical_idx`.
//
// Case 1: if `head_len == 0 && tail_len == 0`
// Everything was drained, reset the head index back to 0.
// H
// [ . . . . . d d d d . . . . . ]
// H
// [ . . . . . . . . . . . . . . ]
//
// Case 2: else if `tail_len == 0`
// Don't move data or the head index.
// H
// [ . . . h h h h d d d d . . . ]
// H
// [ . . . h h h h . . . . . . . ]
//
// Case 3: else if `head_len == 0`
// Don't move data, but move the head index.
// H
// [ . . . d d d d t t t t . . . ]
// H
// [ . . . . . . . t t t t . . . ]
//
// Case 4: else if `tail_len <= head_len`
// Move data, but not the head index.
// H
// [ . . h h h h d d d d t t . . ]
// H
// [ . . h h h h t t . . . . . . ]
//
// Case 5: else
// Move data and the head index.
// H
// [ . . h h d d d d t t t t . . ]
// H
// [ . . . . . . h h t t t t . . ]

// When draining at the front (`.drain(..n)`) or at the back (`.drain(n..)`),
// we don't need to copy any data. The number of elements copied would be 0.
if head_len != 0 && tail_len != 0 {
join_head_and_tail_wrapping(source_deque, drain_len, head_len, tail_len);
// Marking this function as cold helps LLVM to eliminate it entirely if
// this branch is never taken.
// We use `#[cold]` instead of `#[inline(never)]`, because inlining this
// function into the general case (`.drain(n..m)`) is fine.
// See `tests/codegen-llvm/vecdeque-drain.rs` for a test.
#[cold]
fn join_head_and_tail_wrapping<T, A: Allocator>(
source_deque: &mut VecDeque<T, A>,
drain_len: usize,
head_len: usize,
tail_len: usize,
) {
// Pick whether to move the head or the tail here.
let (src, dst, len);
if head_len < tail_len {
src = source_deque.head;
dst = source_deque.to_wrapped_index(drain_len);
len = head_len;
} else {
src = source_deque.to_wrapped_index(head_len + drain_len);
dst = source_deque.to_wrapped_index(head_len);
len = tail_len;
};
// When draining at the front (`.drain(..n)`) or at the back (`.drain(n..)`),
// we don't need to copy any data. The number of elements copied would be 0.
if head_len != 0 && tail_len != 0 {
join_head_and_tail_wrapping(source_deque, drain_len, head_len, tail_len);
// Marking this function as cold helps LLVM to eliminate it entirely if
// this branch is never taken.
// We use `#[cold]` instead of `#[inline(never)]`, because inlining this
// function into the general case (`.drain(n..m)`) is fine.
// See `tests/codegen-llvm/vecdeque-drain.rs` for a test.
#[cold]
fn join_head_and_tail_wrapping<T, A: Allocator>(
source_deque: &mut VecDeque<T, A>,
drain_len: usize,
head_len: usize,
tail_len: usize,
) {
// Pick whether to move the head or the tail here.
let (src, dst, len);
if head_len < tail_len {
src = source_deque.head;
dst = source_deque.to_wrapped_index(drain_len);
len = head_len;
} else {
src = source_deque.to_wrapped_index(head_len + drain_len);
dst = source_deque.to_wrapped_index(head_len);
len = tail_len;
};

// ignore-tidy-undocumented-unsafe
unsafe {
source_deque.wrap_copy(src, dst, len);
// ignore-tidy-undocumented-unsafe
unsafe {
source_deque.wrap_copy(src, dst, len);
}
}
}
}

if new_len == 0 {
// Special case: If the entire deque was drained, reset the head back to 0,
// like `.clear()` does.
source_deque.head = WrappedIndex::zero();
} else if head_len < tail_len {
// If we moved the head above, then we need to adjust the head index here.
source_deque.head = source_deque.to_wrapped_index(drain_len);
if new_len == 0 {
// Special case: If the entire deque was drained, reset the head back to 0,
// like `.clear()` does.
source_deque.head = WrappedIndex::zero();
} else if head_len < tail_len {
// If we moved the head above, then we need to adjust the head index here.
source_deque.head = source_deque.to_wrapped_index(drain_len);
}
source_deque.len = new_len;
}
source_deque.len = new_len;
});

if mem::needs_drop::<T>() && guard.remaining != 0 {
// SAFETY: We just checked that `self.remaining != 0`.
let (front, back) = unsafe { guard.as_slices() };
// since idx is a logical index, we don't need to worry about wrapping.
guard.idx += front.len();
guard.remaining -= front.len();
// SAFETY: This can't have been dropped before since
// `idx` & `remaining` track what's been dropped.
unsafe { ptr::drop_in_place(front) };
guard.remaining = 0;
// SAFETY: Ditto.
unsafe { ptr::drop_in_place(back) };
}
}
}
Expand Down
35 changes: 20 additions & 15 deletions library/alloc/src/slice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -408,30 +408,35 @@ impl<T> [T] {
impl<T: Clone> ConvertVec for T {
#[inline]
default fn to_vec<A: Allocator>(s: &[Self], alloc: A) -> Vec<Self, A> {
use core::mem::DropGuard;

let mut guard = DropGuard::new(
(0, Vec::with_capacity_in(s.len(), alloc)),
|(num_init, mut vec)| {
struct DropGuard<'a, T, A: Allocator> {
vec: &'a mut Vec<T, A>,
num_init: usize,
}
impl<'a, T, A: Allocator> Drop for DropGuard<'a, T, A> {
#[inline]
fn drop(&mut self) {
// SAFETY:
// items were marked initialized in the loop below
unsafe { vec.set_len(num_init) }
},
);
let (num_init, vec) = &mut *guard;

let slots = vec.spare_capacity_mut();
unsafe {
self.vec.set_len(self.num_init);
}
}
}
let mut vec = Vec::with_capacity_in(s.len(), alloc);
let mut guard = DropGuard { vec: &mut vec, num_init: 0 };
let slots = guard.vec.spare_capacity_mut();
// .take(slots.len()) is necessary for LLVM to remove bounds checks
// and has better codegen than zip.
for (i, b) in s.iter().enumerate().take(slots.len()) {
*num_init = i;
guard.num_init = i;
slots[i].write(b.clone());
}

let (_, mut vec) = DropGuard::dismiss(guard);
core::mem::forget(guard);
// SAFETY:
// the vec was allocated and initialized above to at least this length.
unsafe { vec.set_len(s.len()) };
unsafe {
vec.set_len(s.len());
}
vec
}
}
Expand Down
Loading