Rollup merge of #95298 - jhorstmann:fix-double-drop-of-allocator-in-vec-into-iter, r=oli-obk

Fix double drop of allocator in IntoIter impl of Vec

Fixes #95269

The `drop` impl of `IntoIter` reconstructs a `RawVec` from `buf`, `cap` and `alloc`, when that `RawVec` is dropped it also drops the allocator. To avoid dropping the allocator twice we wrap it in `ManuallyDrop` in the `InttoIter` struct.

Note this is my first contribution to the standard library, so I might be missing some details or a better way to solve this.

Author: Dylan-DPC

library/alloc/src/vec/into_iter.rs

@@ -8,7 +8,8 @@ use core::iter::{
     FusedIterator, InPlaceIterable, SourceIter, TrustedLen, TrustedRandomAccessNoCoerce,
 };
 use core::marker::PhantomData;
-use core::mem::{self};
+use core::mem::{self, ManuallyDrop};
+use core::ops::Deref;
 use core::ptr::{self, NonNull};
 use core::slice::{self};
 
@@ -32,7 +33,9 @@ pub struct IntoIter<
     pub(super) buf: NonNull<T>,
     pub(super) phantom: PhantomData<T>,
     pub(super) cap: usize,
-    pub(super) alloc: A,
+    // the drop impl reconstructs a RawVec from buf, cap and alloc
+    // to avoid dropping the allocator twice we need to wrap it into ManuallyDrop
+    pub(super) alloc: ManuallyDrop<A>,
     pub(super) ptr: *const T,
     pub(super) end: *const T,
 }
@@ -295,11 +298,11 @@ where
 impl<T: Clone, A: Allocator + Clone> Clone for IntoIter<T, A> {
     #[cfg(not(test))]
     fn clone(&self) -> Self {
-        self.as_slice().to_vec_in(self.alloc.clone()).into_iter()
+        self.as_slice().to_vec_in(self.alloc.deref().clone()).into_iter()
     }
     #[cfg(test)]
     fn clone(&self) -> Self {
-        crate::slice::to_vec(self.as_slice(), self.alloc.clone()).into_iter()
+        crate::slice::to_vec(self.as_slice(), self.alloc.deref().clone()).into_iter()
     }
 }
 
@@ -311,8 +314,8 @@ 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
-                    let alloc = ptr::read(&self.0.alloc);
+                    // `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_raw_parts_in(self.0.buf.as_ptr(), self.0.cap, alloc);
                 }

library/alloc/src/vec/mod.rs

@@ -2579,7 +2579,7 @@ impl<T, A: Allocator> IntoIterator for Vec<T, A> {
     fn into_iter(self) -> IntoIter<T, A> {
         unsafe {
             let mut me = ManuallyDrop::new(self);
-            let alloc = ptr::read(me.allocator());
+            let alloc = ManuallyDrop::new(ptr::read(me.allocator()));
             let begin = me.as_mut_ptr();
             let end = if mem::size_of::<T>() == 0 {
                 arith_offset(begin as *const i8, me.len() as isize) as *const T

library/alloc/tests/vec.rs

@@ -1,3 +1,6 @@
+use core::alloc::{Allocator, Layout};
+use core::ptr::NonNull;
+use std::alloc::System;
 use std::assert_matches::assert_matches;
 use std::borrow::Cow;
 use std::cell::Cell;
@@ -991,6 +994,31 @@ fn test_into_iter_advance_by() {
     assert_eq!(i.len(), 0);
 }
 
+#[test]
+fn test_into_iter_drop_allocator() {
+    struct ReferenceCountedAllocator<'a>(DropCounter<'a>);
+
+    unsafe impl Allocator for ReferenceCountedAllocator<'_> {
+        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, core::alloc::AllocError> {
+            System.allocate(layout)
+        }
+
+        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+            System.deallocate(ptr, layout)
+        }
+    }
+
+    let mut drop_count = 0;
+
+    let allocator = ReferenceCountedAllocator(DropCounter { count: &mut drop_count });
+    let _ = Vec::<u32, _>::new_in(allocator);
+    assert_eq!(drop_count, 1);
+
+    let allocator = ReferenceCountedAllocator(DropCounter { count: &mut drop_count });
+    let _ = Vec::<u32, _>::new_in(allocator).into_iter();
+    assert_eq!(drop_count, 2);
+}
+
 #[test]
 fn test_from_iter_specialization() {
     let src: Vec<usize> = vec![0usize; 1];