diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index cc6dfb0e33017..03c532bb69769 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/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, pub(super) phantom: PhantomData, 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, pub(super) ptr: *const T, pub(super) end: *const T, } @@ -295,11 +298,11 @@ where impl Clone for IntoIter { #[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 { impl 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); } diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 1ca5ee5537582..9773ec023377f 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -2579,7 +2579,7 @@ impl IntoIterator for Vec { fn into_iter(self) -> IntoIter { 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::() == 0 { arith_offset(begin as *const i8, me.len() as isize) as *const T diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index ca0fcc855c7b8..19e39ebf910b5 100644 --- a/library/alloc/tests/vec.rs +++ b/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, core::alloc::AllocError> { + System.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + System.deallocate(ptr, layout) + } + } + + let mut drop_count = 0; + + let allocator = ReferenceCountedAllocator(DropCounter { count: &mut drop_count }); + let _ = Vec::::new_in(allocator); + assert_eq!(drop_count, 1); + + let allocator = ReferenceCountedAllocator(DropCounter { count: &mut drop_count }); + let _ = Vec::::new_in(allocator).into_iter(); + assert_eq!(drop_count, 2); +} + #[test] fn test_from_iter_specialization() { let src: Vec = vec![0usize; 1];