From ce994027fe2ce57a4759a4c1327dbfb3592bfdff Mon Sep 17 00:00:00 2001 From: The8472 Date: Mon, 10 May 2021 21:49:53 +0200 Subject: [PATCH 1/3] replace vec::Drain drop loops with drop_in_place --- library/alloc/src/vec/drain.rs | 51 ++++++++++++++++++++++++---------- 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/library/alloc/src/vec/drain.rs b/library/alloc/src/vec/drain.rs index ff98091a0d2ab..75aa35c926ca0 100644 --- a/library/alloc/src/vec/drain.rs +++ b/library/alloc/src/vec/drain.rs @@ -1,7 +1,7 @@ use crate::alloc::{Allocator, Global}; use core::fmt; use core::iter::{FusedIterator, TrustedLen}; -use core::mem::{self}; +use core::mem::{self, MaybeUninit}; use core::ptr::{self, NonNull}; use core::slice::{self}; @@ -104,16 +104,11 @@ impl DoubleEndedIterator for Drain<'_, T, A> { #[stable(feature = "drain", since = "1.6.0")] impl 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(); @@ -131,15 +126,43 @@ impl 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::() == 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 { + let vec = vec.as_mut(); + let spare_capacity = vec.spare_capacity_mut(); + let drop_offset = drop_ptr.offset_from(spare_capacity.as_ptr() as *const _) as usize; + let drop_range = drop_offset..(drop_offset + drop_len); + let to_drop = &mut spare_capacity[drop_range]; + ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(to_drop)); + } } } From 6851b8d931381c4e7b6d60dab70002efd85dc7a9 Mon Sep 17 00:00:00 2001 From: The8472 Date: Sat, 20 Nov 2021 01:29:04 +0100 Subject: [PATCH 2/3] document why we're not directly passing drop_ptr to drop_in_place --- library/alloc/src/vec/drain.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/library/alloc/src/vec/drain.rs b/library/alloc/src/vec/drain.rs index 75aa35c926ca0..053ee8033ba92 100644 --- a/library/alloc/src/vec/drain.rs +++ b/library/alloc/src/vec/drain.rs @@ -156,6 +156,10 @@ impl Drop for Drain<'_, T, A> { } 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 = vec.as_mut(); let spare_capacity = vec.spare_capacity_mut(); let drop_offset = drop_ptr.offset_from(spare_capacity.as_ptr() as *const _) as usize; From 2d8a11bdbb2623a4e2983870006cbd00eb210ffb Mon Sep 17 00:00:00 2001 From: The 8472 Date: Thu, 9 Dec 2021 00:20:13 +0100 Subject: [PATCH 3/3] Use `*mut [T]` instead of `[MaybeUninit]` --- library/alloc/src/vec/drain.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/library/alloc/src/vec/drain.rs b/library/alloc/src/vec/drain.rs index 053ee8033ba92..089cd4db64c1d 100644 --- a/library/alloc/src/vec/drain.rs +++ b/library/alloc/src/vec/drain.rs @@ -1,7 +1,7 @@ use crate::alloc::{Allocator, Global}; use core::fmt; use core::iter::{FusedIterator, TrustedLen}; -use core::mem::{self, MaybeUninit}; +use core::mem; use core::ptr::{self, NonNull}; use core::slice::{self}; @@ -160,12 +160,10 @@ impl Drop for Drain<'_, T, A> { // 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 = vec.as_mut(); - let spare_capacity = vec.spare_capacity_mut(); - let drop_offset = drop_ptr.offset_from(spare_capacity.as_ptr() as *const _) as usize; - let drop_range = drop_offset..(drop_offset + drop_len); - let to_drop = &mut spare_capacity[drop_range]; - ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(to_drop)); + 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); } } }