From e883282c80b71c9c4508baa551301965433c60ba Mon Sep 17 00:00:00 2001
From: soruh <mail@soruh.de>
Date: Fri, 31 Jul 2026 04:04:46 +0200
Subject: [PATCH] Retain cleared allocations during unwinding

---
 src/lib.rs    |  5 +++++
 src/tests.rs  | 48 ++++++++++++++++++++++++++++++++++++++++++------
 src/vtable.rs | 36 ++++++++++++++++++++----------------
 3 files changed, 67 insertions(+), 22 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 231bf81..4ee7e02 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -88,6 +88,11 @@ impl TypeErasedVec {
     }
 
     /// Clear any remaining elements in the vector, dropping them.
+    ///
+    /// # Panics
+    ///
+    /// If an element destructor panics, the vector remains valid and empty and
+    /// retains its allocation.
     pub fn clear(&mut self) {
         // SAFETY: The clear function correctly targets the currently stored type elements.
         unsafe { (self.vtable.clear)(&mut self.parts) };
diff --git a/src/tests.rs b/src/tests.rs
index 62bdca3..ae50fc9 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -323,23 +323,48 @@ fn test_reserve_is_unwind_safe() {
 
 #[test]
 fn test_clear_is_unwind_safe_when_element_drop_panics() {
-    struct PanicOnFirstDrop {
+    struct PanicOnSelectedDrop {
+        drop_count: Rc<Cell<usize>>,
         has_panicked: Rc<Cell<bool>>,
+        should_panic: bool,
     }
 
-    impl Drop for PanicOnFirstDrop {
+    impl Drop for PanicOnSelectedDrop {
         fn drop(&mut self) {
+            let Some(next_count) = self.drop_count.get().checked_add(1) else {
+                panic!("drop counter overflow");
+            };
+            self.drop_count.set(next_count);
+
             assert!(
-                self.has_panicked.replace(true),
+                !self.should_panic || self.has_panicked.replace(true),
                 "simulated panic while clearing an element"
             );
         }
     }
 
+    let drop_count = Rc::new(Cell::new(0));
     let has_panicked = Rc::new(Cell::new(false));
-    let mut erased = TypeErasedVec::new(vec![PanicOnFirstDrop {
-        has_panicked: Rc::clone(&has_panicked),
-    }]);
+    let values = vec![
+        PanicOnSelectedDrop {
+            drop_count: Rc::clone(&drop_count),
+            has_panicked: Rc::clone(&has_panicked),
+            should_panic: false,
+        },
+        PanicOnSelectedDrop {
+            drop_count: Rc::clone(&drop_count),
+            has_panicked: Rc::clone(&has_panicked),
+            should_panic: true,
+        },
+        PanicOnSelectedDrop {
+            drop_count: Rc::clone(&drop_count),
+            has_panicked: Rc::clone(&has_panicked),
+            should_panic: false,
+        },
+    ];
+    let original_ptr = values.as_ptr();
+    let original_capacity = values.capacity();
+    let mut erased = TypeErasedVec::new(values);
 
     let result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
         erased.clear();
@@ -347,7 +372,18 @@ fn test_clear_is_unwind_safe_when_element_drop_panics() {
 
     assert!(result.is_err());
     assert!(has_panicked.get());
+    assert_eq!(drop_count.get(), 3);
     assert!(erased.is_empty());
+    assert_eq!(erased.capacity(), original_capacity);
+
+    {
+        // SAFETY: Clear retains the allocation's original element type.
+        let guard = unsafe { erased.cast_type::<PanicOnSelectedDrop>() };
+        assert_eq!(guard.as_slice().as_ptr(), original_ptr);
+    }
+
+    drop(erased);
+    assert_eq!(drop_count.get(), 3);
 }
 
 #[test]
diff --git a/src/vtable.rs b/src/vtable.rs
index 0fb127a..1cf48a5 100644
--- a/src/vtable.rs
+++ b/src/vtable.rs
@@ -33,6 +33,14 @@ impl<'parts, T> VecPartsRestoreGuard<'parts, T> {
     fn length(&self) -> usize {
         self.vec.len()
     }
+
+    fn capacity(&self) -> usize {
+        self.vec.capacity()
+    }
+
+    fn pointer(&self) -> *const T {
+        self.vec.as_ptr()
+    }
 }
 
 impl<T> Drop for VecPartsRestoreGuard<'_, T> {
@@ -61,23 +69,19 @@ impl TypeErasedVecVtable {
         }
 
         unsafe fn clear_vec<T>(parts: &mut VecParts) {
-            let original_parts = std::mem::replace(parts, VecParts::from_vec(Vec::<T>::new()));
+            let original_ptr = parts.ptr.as_ptr();
+            let original_capacity = parts.cap;
 
-            // SAFETY: We transferred ownership of the original allocation out of
-            // `parts` before reconstructing the Vec. If `clear` unwinds, the Vec
-            // owns and frees that allocation while `parts` remains a valid empty Vec.
-            unsafe {
-                let mut vec = original_parts.into_vec::<T>();
-                vec.clear();
-                let new_parts = VecParts::from_vec(vec);
-                debug_assert!(std::ptr::eq(
-                    new_parts.ptr.as_ptr(),
-                    original_parts.ptr.as_ptr()
-                ));
-                debug_assert_eq!(new_parts.cap, original_parts.cap);
-                debug_assert_eq!(new_parts.len, 0);
-                *parts = new_parts;
-            }
+            // SAFETY: This vtable function is selected only for raw parts that
+            // were produced by Vec<T>.
+            let mut restore_guard = unsafe { VecPartsRestoreGuard::<T>::new(parts) };
+            restore_guard.vec_mut().clear();
+            debug_assert!(std::ptr::eq(
+                restore_guard.pointer().cast::<u8>(),
+                original_ptr
+            ));
+            debug_assert_eq!(restore_guard.capacity(), original_capacity);
+            debug_assert_eq!(restore_guard.length(), 0);
         }
 
         unsafe fn reserve_vec<T>(parts: &mut VecParts, additional: usize) {