Skip to content
Open
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
4 changes: 2 additions & 2 deletions library/alloc/src/collections/vec_deque/spec_extend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,15 +78,15 @@ where

#[cfg(not(test))]
impl<T, A: Allocator> SpecExtend<T, vec::IntoIter<T>> for VecDeque<T, A> {
fn spec_extend(&mut self, mut iterator: vec::IntoIter<T>) {
fn spec_extend(&mut self, iterator: vec::IntoIter<T>) {
let slice = iterator.as_slice();
self.reserve(slice.len());

unsafe {
self.copy_slice(self.to_physical_idx(self.len), slice);
self.len += slice.len();
}
iterator.forget_remaining_elements();
iterator.forget_remaining_elements_and_dealloc();
}
}

Expand Down
50 changes: 46 additions & 4 deletions library/alloc/src/vec/into_iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,12 +160,51 @@ impl<T, A: Allocator> IntoIter<T, A> {
}

/// Forgets to Drop the remaining elements while still allowing the backing allocation to be freed.
///
/// This method does not consume `self`, and leaves deallocation to `impl Drop for IntoIter`.
/// If consuming `self` is possible, consider calling
/// [`Self::forget_remaining_elements_and_dealloc()`] instead.
pub(crate) fn forget_remaining_elements(&mut self) {
// For the ZST case, it is crucial that we mutate `end` here, not `ptr`.
// `ptr` must stay aligned, while `end` may be unaligned.
self.end = self.ptr.as_ptr();
}

/// Forgets to Drop the remaining elements and frees the backing allocation.
/// Consuming version of [`Self::forget_remaining_elements()`].
///
/// This can be used in place of `drop(self)` when `self` is known to be exhausted,
/// to avoid producing a needless `drop_in_place::<[T]>()`.
#[inline]
pub(crate) fn forget_remaining_elements_and_dealloc(self) {
let mut this = ManuallyDrop::new(self);
// SAFETY: `this` is in ManuallyDrop, so it will not be double-freed.
unsafe {
this.dealloc_only();
}
}

/// Frees the allocation, without checking or dropping anything else.
///
/// The safe version of this method is [`Self::forget_remaining_elements_and_dealloc()`].
/// This function exists only to share code between that method and the `impl Drop`.
///
/// # Safety
///
/// This function must only be called with an [`IntoIter`] that is not going to be dropped
/// or otherwise used in any way, either because it is being forgotten or because its `Drop`
/// is already executing; otherwise a double-free will occur, and possibly a read from freed
/// memory if there are any remaining elements.
#[inline]
unsafe fn dealloc_only(&mut self) {
unsafe {
// SAFETY: our caller promises not to touch `*self` again
let alloc = ManuallyDrop::take(&mut self.alloc);
// RawVec handles deallocation
let _ = RawVec::from_nonnull_in(self.buf, self.cap, alloc);
}
}

#[cfg(not(no_global_oom_handling))]
#[inline]
pub(crate) fn into_vecdeque(self) -> VecDeque<T, A> {
Expand Down Expand Up @@ -329,6 +368,12 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
accum = f(accum, tmp);
}
}

// There are in fact no remaining elements to forget, but by doing this we can avoid
// potentially generating a needless loop to drop the elements that cannot exist at
// this point.
self.forget_remaining_elements_and_dealloc();

accum
}

Expand Down Expand Up @@ -496,10 +541,7 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for IntoIter<T, A> {
impl<T, A: Allocator> Drop for DropGuard<'_, T, A> {
fn drop(&mut self) {
unsafe {
// `IntoIter::alloc` is not used anymore after this and will be dropped by RawVec
let alloc = ManuallyDrop::take(&mut self.0.alloc);
// RawVec handles deallocation
let _ = RawVec::from_nonnull_in(self.0.buf, self.0.cap, alloc);
self.0.dealloc_only();
}
}
}
Expand Down
67 changes: 67 additions & 0 deletions tests/codegen-llvm/vec-into-iter-drops.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
//@ compile-flags: -Copt-level=3
// Test that we can avoid generating an element dropping loop when `vec::IntoIter` is consumed.
#![crate_type = "lib"]

use std::vec;

struct Bomb;
impl Drop for Bomb {
#[inline]
fn drop(&mut self) {
panic!("dropped")
}
}

/// Test case originally from https://users.rust-lang.org/t/unnecessary-drop-in-place-emitted-for-a-fully-consumed-intoiter/135119
///
/// What we are looking for is that there should be no calls to `impl Drop for Bomb`
/// because every element is unconditionally forgotten.
//
// CHECK-LABEL: @vec_for_each_doesnt_drop
#[no_mangle]
pub fn vec_for_each_doesnt_drop(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
// CHECK-NOT: panic
// CHECK-NOT: {{call.*drop_in_place}}
// CHECK-NOT: Bomb$u20$as$u20$core..ops..drop..Drop
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: CHECK-NOT tests are generally quite fragile, being prone to accidentally checking for something that never existed at all, or having the implementation in the standard library change such that a previously-useful test no longer does anything useful.

Thus I would suggest that whatever you're checking for here you also write another function that intentionally does trigger whatever you expect to not be there in this one, with positive CHECKs for the same things that are CHECK-NOTs here.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. However, I encounter a problem: I add more tests to this file and LLVM decides that they should call drop_in_place::<IntoIter>() instead of inlining it, defeating the entire point of this change. That’s a sign that this optimization is fragile and more work is needed, I guess. (Or maybe the perf results will show that it’s good despite that.)

I am now trying to write more explicit code that doesn’t rely on inlining to help.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I’ve pushed a new commit 5fa286d that performs explicit forgetting instead of relying on inlining and dead code elimination to achieve anything.

let mut last = 0;
v.into_iter().for_each(|(x, bomb)| {
last = x;
std::mem::forget(bomb);
});
last
}

/// Test that does *not* get the above optimization we are expecting:
/// it uses a normal `for` loop which calls `Iterator::next()` and then drops the iterator,
/// and dropping the iterator drops remaining items.
///
/// This test exists to prove that the above CHECK-NOT is looking for the right things.
/// However, it might start failing if LLVM figures out that there are no remaining items.
//
// CHECK-LABEL: @vec_for_loop
#[no_mangle]
pub fn vec_for_loop(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
// CHECK: {{call.*drop_in_place}}
let mut last = 0;
for (x, bomb) in v {
last = x;
std::mem::forget(bomb);
}
last
}

/// Test where there still should be drops because there are no forgets.
///
/// This test exists to prove that the above CHECK-NOT is looking for the right things
/// and does not say anything interesting about codegen itself.
//
// CHECK-LABEL: @vec_for_each_does_drop
#[no_mangle]
pub fn vec_for_each_does_drop(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
// CHECK: begin_panic
let mut last = 0;
v.into_iter().for_each(|(x, bomb)| {
last = x;
});
last
}
Loading