Skip to content

Commit

Permalink
Auto merge of rust-lang#85157 - the8472:drain-drop-in-place, r=Mark-S…
Browse files Browse the repository at this point in the history
…imulacrum

replace vec::Drain drop loops with drop_in_place

The `Drain::drop` implementation came up in rust-lang#82185 (comment) as potentially interfering with other optimization work due its widespread use somewhere in `println!`

`@rustbot` label T-libs-impl
  • Loading branch information
bors committed Dec 9, 2021
2 parents 3b263ce + 2d8a11b commit 0b42dea
Showing 1 changed file with 39 additions and 14 deletions.
53 changes: 39 additions & 14 deletions library/alloc/src/vec/drain.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::alloc::{Allocator, Global};
use core::fmt;
use core::iter::{FusedIterator, TrustedLen};
use core::mem::{self};
use core::mem;
use core::ptr::{self, NonNull};
use core::slice::{self};

Expand Down Expand Up @@ -104,16 +104,11 @@ impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
#[stable(feature = "drain", since = "1.6.0")]
impl<T, A: Allocator> Drop for Drain<'_, T, A> {
fn drop(&mut self) {
/// Continues dropping the remaining elements in the `Drain`, then moves back the
/// un-`Drain`ed elements to restore the original `Vec`.
/// Moves back the un-`Drain`ed elements to restore the original `Vec`.
struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>);

impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
fn drop(&mut self) {
// Continue the same loop we have below. If the loop already finished, this does
// nothing.
self.0.for_each(drop);

if self.0.tail_len > 0 {
unsafe {
let source_vec = self.0.vec.as_mut();
Expand All @@ -131,15 +126,45 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
}
}

// exhaust self first
while let Some(item) = self.next() {
let guard = DropGuard(self);
drop(item);
mem::forget(guard);
let iter = mem::replace(&mut self.iter, (&mut []).iter());
let drop_len = iter.len();
let drop_ptr = iter.as_slice().as_ptr();

// forget iter so there's no aliasing reference
drop(iter);

let mut vec = self.vec;

if mem::size_of::<T>() == 0 {
// ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount.
// this can be achieved by manipulating the Vec length instead of moving values out from `iter`.
unsafe {
let vec = vec.as_mut();
let old_len = vec.len();
vec.set_len(old_len + drop_len + self.tail_len);
vec.truncate(old_len + self.tail_len);
}

return;
}

// ensure elements are moved back into their appropriate places, even when drop_in_place panics
let _guard = DropGuard(self);

if drop_len == 0 {
return;
}

// Drop a `DropGuard` to move back the non-drained tail of `self`.
DropGuard(self);
unsafe {
// drop_ptr comes from a slice::Iter which only gives us a &[T] but for drop_in_place
// a pointer with mutable provenance is necessary. Therefore we must reconstruct
// it from the original vec but also avoid creating a &mut to the front since that could
// invalidate raw pointers to it which some unsafe code might rely on.
let vec_ptr = vec.as_mut().as_mut_ptr();
let drop_offset = drop_ptr.offset_from(vec_ptr) as usize;
let to_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len);
ptr::drop_in_place(to_drop);
}
}
}

Expand Down

0 comments on commit 0b42dea

Please sign in to comment.