Skip to content

Commit 5fa286d

Browse files
committed
Explicitly deallocate instead of dropping the IntoIter.
This has the same end goal as the previous commit but does not rely on compiler optimizations to delete the unwanted code; instead it enters a code path which explicitly frees the allocation and forgets the `IntoIter`.
1 parent 8052601 commit 5fa286d

File tree

3 files changed

+86
-11
lines changed

3 files changed

+86
-11
lines changed

library/alloc/src/collections/vec_deque/spec_extend.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,15 +78,15 @@ where
7878

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

8585
unsafe {
8686
self.copy_slice(self.to_physical_idx(self.len), slice);
8787
self.len += slice.len();
8888
}
89-
iterator.forget_remaining_elements();
89+
iterator.forget_remaining_elements_and_dealloc();
9090
}
9191
}
9292

library/alloc/src/vec/into_iter.rs

Lines changed: 41 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -160,12 +160,51 @@ impl<T, A: Allocator> IntoIter<T, A> {
160160
}
161161

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

173+
/// Forgets to Drop the remaining elements and frees the backing allocation.
174+
/// Consuming version of [`Self::forget_remaining_elements()`].
175+
///
176+
/// This can be used in place of `drop(self)` when `self` is known to be exhausted,
177+
/// to avoid producing a needless `drop_in_place::<[T]>()`.
178+
#[inline]
179+
pub(crate) fn forget_remaining_elements_and_dealloc(self) {
180+
let mut this = ManuallyDrop::new(self);
181+
// SAFETY: `this` is in ManuallyDrop, so it will not be double-freed.
182+
unsafe {
183+
this.dealloc_only();
184+
}
185+
}
186+
187+
/// Frees the allocation, without checking or dropping anything else.
188+
///
189+
/// The safe version of this method is [`Self::forget_remaining_elements_and_dealloc()`].
190+
/// This function exists only to share code between that method and the `impl Drop`.
191+
///
192+
/// # Safety
193+
///
194+
/// This function must only be called with an [`IntoIter`] that is not going to be dropped
195+
/// or otherwise used in any way, either because it is being forgotten or because its `Drop`
196+
/// is already executing; otherwise a double-free will occur, and possibly a read from freed
197+
/// memory if there are any remaining elements.
198+
#[inline]
199+
unsafe fn dealloc_only(&mut self) {
200+
unsafe {
201+
// SAFETY: our caller promises not to touch `*self` again
202+
let alloc = ManuallyDrop::take(&mut self.alloc);
203+
// RawVec handles deallocation
204+
let _ = RawVec::from_nonnull_in(self.buf, self.cap, alloc);
205+
}
206+
}
207+
169208
#[cfg(not(no_global_oom_handling))]
170209
#[inline]
171210
pub(crate) fn into_vecdeque(self) -> VecDeque<T, A> {
@@ -333,7 +372,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
333372
// There are in fact no remaining elements to forget, but by doing this we can avoid
334373
// potentially generating a needless loop to drop the elements that cannot exist at
335374
// this point.
336-
self.forget_remaining_elements();
375+
self.forget_remaining_elements_and_dealloc();
337376

338377
accum
339378
}
@@ -502,10 +541,7 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for IntoIter<T, A> {
502541
impl<T, A: Allocator> Drop for DropGuard<'_, T, A> {
503542
fn drop(&mut self) {
504543
unsafe {
505-
// `IntoIter::alloc` is not used anymore after this and will be dropped by RawVec
506-
let alloc = ManuallyDrop::take(&mut self.0.alloc);
507-
// RawVec handles deallocation
508-
let _ = RawVec::from_nonnull_in(self.0.buf, self.0.cap, alloc);
544+
self.0.dealloc_only();
509545
}
510546
}
511547
}

tests/codegen-llvm/vec-into-iter-drops.rs

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,16 @@ impl Drop for Bomb {
1212
}
1313
}
1414

15-
/// This test case from https://users.rust-lang.org/t/unnecessary-drop-in-place-emitted-for-a-fully-consumed-intoiter/135119
16-
/// It should not emit any `drop::<Bomb>()` because every element is forgotten.
17-
// CHECK-LABEL: @vec_into_iter_drop_option
15+
/// Test case originally from https://users.rust-lang.org/t/unnecessary-drop-in-place-emitted-for-a-fully-consumed-intoiter/135119
16+
///
17+
/// What we are looking for is that there should be no calls to `impl Drop for Bomb`
18+
/// because every element is unconditionally forgotten.
19+
//
20+
// CHECK-LABEL: @vec_for_each_doesnt_drop
1821
#[no_mangle]
19-
pub fn vec_into_iter_drop_option(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
22+
pub fn vec_for_each_doesnt_drop(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
2023
// CHECK-NOT: panic
24+
// CHECK-NOT: {{call.*drop_in_place}}
2125
// CHECK-NOT: Bomb$u20$as$u20$core..ops..drop..Drop
2226
let mut last = 0;
2327
v.into_iter().for_each(|(x, bomb)| {
@@ -26,3 +30,38 @@ pub fn vec_into_iter_drop_option(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
2630
});
2731
last
2832
}
33+
34+
/// Test that does *not* get the above optimization we are expecting:
35+
/// it uses a normal `for` loop which calls `Iterator::next()` and then drops the iterator,
36+
/// and dropping the iterator drops remaining items.
37+
///
38+
/// This test exists to prove that the above CHECK-NOT is looking for the right things.
39+
/// However, it might start failing if LLVM figures out that there are no remaining items.
40+
//
41+
// CHECK-LABEL: @vec_for_loop
42+
#[no_mangle]
43+
pub fn vec_for_loop(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
44+
// CHECK: {{call.*drop_in_place}}
45+
let mut last = 0;
46+
for (x, bomb) in v {
47+
last = x;
48+
std::mem::forget(bomb);
49+
}
50+
last
51+
}
52+
53+
/// Test where there still should be drops because there are no forgets.
54+
///
55+
/// This test exists to prove that the above CHECK-NOT is looking for the right things
56+
/// and does not say anything interesting about codegen itself.
57+
//
58+
// CHECK-LABEL: @vec_for_each_does_drop
59+
#[no_mangle]
60+
pub fn vec_for_each_does_drop(v: vec::Vec<(usize, Option<Bomb>)>) -> usize {
61+
// CHECK: begin_panic
62+
let mut last = 0;
63+
v.into_iter().for_each(|(x, bomb)| {
64+
last = x;
65+
});
66+
last
67+
}

0 commit comments

Comments
 (0)