fix: Drop guard was deallocating with the incorrect size

InPlaceDstBufDrop holds onto the allocation before the shrinking happens
which means it must deallocate the destination elements but the source
allocation.

3 files changed, 25 insertions, 13 deletions

diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs
index ec5f32539f2..633ca8c99b4 100644
--- a/library/alloc/src/vec/in_place_collect.rs
+++ b/library/alloc/src/vec/in_place_collect.rs
@@ -72,7 +72,7 @@
 //! This is handled by the [`InPlaceDrop`] guard for sink items (`U`) and by
 //! [`vec::IntoIter::forget_allocation_drop_remaining()`] for remaining source items (`T`).
 //!
-//! If dropping any remaining source item (`T`) panics then [`InPlaceDstBufDrop`] will handle dropping
+//! If dropping any remaining source item (`T`) panics then [`InPlaceDstDataSrcBufDrop`] will handle dropping
 //! the already collected sink items (`U`) and freeing the allocation.
 //!
 //! [`vec::IntoIter::forget_allocation_drop_remaining()`]: super::IntoIter::forget_allocation_drop_remaining()
@@ -158,11 +158,12 @@ use crate::alloc::{handle_alloc_error, Global};
 use core::alloc::Allocator;
 use core::alloc::Layout;
 use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
+use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop, SizedTypeProperties};
 use core::num::NonZeroUsize;
 use core::ptr::{self, NonNull};
 
-use super::{InPlaceDrop, InPlaceDstBufDrop, SpecFromIter, SpecFromIterNested, Vec};
+use super::{InPlaceDrop, InPlaceDstDataSrcBufDrop, SpecFromIter, SpecFromIterNested, Vec};
 
 const fn in_place_collectible<DEST, SRC>(
     step_merge: Option<NonZeroUsize>,
@@ -262,7 +263,7 @@ where
             );
         }
 
-        // The ownership of the allocation and the new `T` values is temporarily moved into `dst_guard`.
+        // The ownership of the source allocation and the new `T` values is temporarily moved into `dst_guard`.
         // This is safe because
         // * `forget_allocation_drop_remaining` immediately forgets the allocation
         // before any panic can occur in order to avoid any double free, and then proceeds to drop
@@ -273,7 +274,8 @@ where
         // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce
         // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the
         // module documentation why this is ok anyway.
-        let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap: dst_cap };
+        let dst_guard =
+            InPlaceDstDataSrcBufDrop { ptr: dst_buf, len, src_cap, src: PhantomData::<I::Src> };
         src.forget_allocation_drop_remaining();
 
         // Adjust the allocation if the alignment didn't match or the source had a capacity in bytes
diff --git a/library/alloc/src/vec/in_place_drop.rs b/library/alloc/src/vec/in_place_drop.rs
index 25ca33c6a7b..40a540b57fc 100644
--- a/library/alloc/src/vec/in_place_drop.rs
+++ b/library/alloc/src/vec/in_place_drop.rs
@@ -1,6 +1,10 @@
-use core::ptr::{self};
+use core::marker::PhantomData;
+use core::ptr::{self, drop_in_place};
 use core::slice::{self};
 
+use crate::alloc::Global;
+use crate::raw_vec::RawVec;
+
 // A helper struct for in-place iteration that drops the destination slice of iteration,
 // i.e. the head. The source slice (the tail) is dropped by IntoIter.
 pub(super) struct InPlaceDrop<T> {
@@ -23,17 +27,23 @@ impl<T> Drop for InPlaceDrop<T> {
     }
 }
 
-// A helper struct for in-place collection that drops the destination allocation and elements,
-// to avoid leaking them if some other destructor panics.
-pub(super) struct InPlaceDstBufDrop<T> {
-    pub(super) ptr: *mut T,
+// A helper struct for in-place collection that drops the destination items together with
+// the source allocation - i.e. before the reallocation happened - to avoid leaking them
+// if some other destructor panics.
+pub(super) struct InPlaceDstDataSrcBufDrop<Src, Dest> {
+    pub(super) ptr: *mut Dest,
     pub(super) len: usize,
-    pub(super) cap: usize,
+    pub(super) src_cap: usize,
+    pub(super) src: PhantomData<Src>,
 }
 
-impl<T> Drop for InPlaceDstBufDrop<T> {
+impl<Src, Dest> Drop for InPlaceDstDataSrcBufDrop<Src, Dest> {
     #[inline]
     fn drop(&mut self) {
-        unsafe { super::Vec::from_raw_parts(self.ptr, self.len, self.cap) };
+        unsafe {
+            let _drop_allocation =
+                RawVec::<Src>::from_raw_parts_in(self.ptr.cast::<Src>(), self.src_cap, Global);
+            drop_in_place(core::ptr::slice_from_raw_parts_mut::<Dest>(self.ptr, self.len));
+        };
     }
 }
diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index 35ea97bfe60..48daec15228 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -123,7 +123,7 @@ use self::set_len_on_drop::SetLenOnDrop;
 mod set_len_on_drop;
 
 #[cfg(not(no_global_oom_handling))]
-use self::in_place_drop::{InPlaceDrop, InPlaceDstBufDrop};
+use self::in_place_drop::{InPlaceDrop, InPlaceDstDataSrcBufDrop};
 
 #[cfg(not(no_global_oom_handling))]
 mod in_place_drop;