Auto merge of #148486 - kpreid:vec-iter-drop, r=jhpratt

Explicitly forget the zero remaining elements in `vec::IntoIter::fold()`.



[Original description:] ~~This seems to help LLVM notice that dropping the elements in the destructor of `IntoIter` is not necessary. In cases it doesn’t help, it should be cheap since it is just one assignment.~~

This PR adds a function to `vec::IntoIter()` which is used used by `fold()` and `spec_extend()`, when those operations complete, to forget the zero remaining elements and only deallocate the allocation, ensuring that there will never be a useless loop to drop zero remaining elements when the iterator is dropped.

This is my first ever attempt at this kind of codegen micro-optimization in the standard library, so please let me know what should go into the PR or what sort of additional systematic testing might indicate this is a good or bad idea.

Author: bors

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

@@ -78,15 +78,15 @@ where
 
 #[cfg(not(test))]
 impl<T, A1: Allocator, A2: Allocator> SpecExtend<T, vec::IntoIter<T, A2>> for VecDeque<T, A1> {
-    fn spec_extend(&mut self, mut iterator: vec::IntoIter<T, A2>) {
+    fn spec_extend(&mut self, iterator: vec::IntoIter<T, A2>) {
         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();
     }
 }
 
@@ -155,12 +155,12 @@ where
 #[cfg(not(test))]
 impl<T, A1: Allocator, A2: Allocator> SpecExtendFront<T, vec::IntoIter<T, A2>> for VecDeque<T, A1> {
     #[track_caller]
-    fn spec_extend_front(&mut self, mut iterator: vec::IntoIter<T, A2>) {
+    fn spec_extend_front(&mut self, iterator: vec::IntoIter<T, A2>) {
         let slice = iterator.as_slice();
         self.reserve(slice.len());
         // SAFETY: `slice.len()` space was just reserved and elements in the slice are forgotten after this call
         unsafe { prepend_reversed(self, slice) };
-        iterator.forget_remaining_elements();
+        iterator.forget_remaining_elements_and_dealloc();
     }
 }
 
@@ -170,12 +170,12 @@ impl<T, A1: Allocator, A2: Allocator> SpecExtendFront<T, Rev<vec::IntoIter<T, A2
 {
     #[track_caller]
     fn spec_extend_front(&mut self, iterator: Rev<vec::IntoIter<T, A2>>) {
-        let mut iterator = iterator.into_inner();
+        let iterator = iterator.into_inner();
         let slice = iterator.as_slice();
         self.reserve(slice.len());
         // SAFETY: `slice.len()` space was just reserved and elements in the slice are forgotten after this call
         unsafe { prepend(self, slice) };
-        iterator.forget_remaining_elements();
+        iterator.forget_remaining_elements_and_dealloc();
     }
 }
 

library/alloc/src/vec/into_iter.rs

@@ -159,12 +159,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> {
@@ -328,6 +367,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
     }
 
@@ -500,10 +545,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();
                 }
             }
         }

library/alloc/src/vec/spec_extend.rs

@@ -29,11 +29,11 @@ where
 }
 
 impl<T, A1: Allocator, A2: Allocator> SpecExtend<T, IntoIter<T, A2>> for Vec<T, A1> {
-    fn spec_extend(&mut self, mut iterator: IntoIter<T, A2>) {
+    fn spec_extend(&mut self, iterator: IntoIter<T, A2>) {
         unsafe {
             self.append_elements(iterator.as_slice() as _);
         }
-        iterator.forget_remaining_elements();
+        iterator.forget_remaining_elements_and_dealloc();
     }
 }
 

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

@@ -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: drop_in_place
+    // CHECK-NOT: Bomb$u20$as$u20$core..ops..drop..Drop
+    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: 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
+}