Restore guarded vectors during unwinding

src/guard.rs

@@ -12,6 +12,33 @@ pub struct ContentGuard<'vec, T> {
     _phantom: PhantomData<T>,
 }
 
+/// Temporarily owns a typed vector while guaranteeing that its latest state is
+/// transferred back to the erased owner from [`Drop`].
+///
+/// This guard makes [`ContentGuard::with`] unwind-safe without rolling back
+/// mutations performed before a panic.
+struct VecRestoreGuard<'vec, T> {
+    erased: &'vec mut TypeErasedVec,
+    vec: ManuallyDrop<Vec<T>>,
+    reerase: fn(Vec<T>) -> TypeErasedVec,
+}
+
+impl<T> VecRestoreGuard<'_, T> {
+    fn vec_mut(&mut self) -> &mut Vec<T> {
+        &mut self.vec
+    }
+}
+
+impl<T> Drop for VecRestoreGuard<'_, T> {
+    fn drop(&mut self) {
+        // SAFETY: `vec` is taken exactly once by this Drop implementation. It
+        // is immediately transferred back into the erased owner without
+        // allocating or dropping any elements.
+        let vec = unsafe { ManuallyDrop::take(&mut self.vec) };
+        *self.erased = (self.reerase)(vec);
+    }
+}
+
 impl<T> fmt::Debug for ContentGuard<'_, T> {
     fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
         formatter
@@ -47,14 +74,19 @@ impl<'vec, T> ContentGuard<'vec, T> {
     /// Calls `function` with the underlying vector and then type-erases it again.
     ///
     /// Mutations made by `function` are retained on normal return.
+    ///
+    /// If `function` panics, the vector's latest state is restored to the
+    /// erased owner during unwinding. Mutations made before the panic are
+    /// preserved rather than rolled back.
     pub fn with<R>(&mut self, f: impl FnOnce(&mut Vec<T>) -> R) -> R {
-        let mut vec = self.take();
-        // This is unwind-safe. If the closure panics, the inner `Vec<T>` drops normally.
-        // The `erased` reference was swapped with a 0-capacity vector inside `take()`,
-        // preventing any double-free or memory leak.
-        let res = f(&mut vec);
-        *self.erased = (self.reerase)(vec);
-        res
+        let vec = self.take();
+        let mut restore_guard = VecRestoreGuard {
+            erased: self.erased,
+            vec: ManuallyDrop::new(vec),
+            reerase: self.reerase,
+        };
+
+        f(restore_guard.vec_mut())
     }
 
     /// Removes and drops all initialized elements.

src/tests.rs

@@ -454,6 +454,8 @@ fn test_guard_with_unwind_safety() {
     let drop_count = Rc::new(Cell::new(0));
     let mut vec = Vec::with_capacity(10);
     vec.push(DropTracker::new(Rc::clone(&drop_count)));
+    let original_ptr = vec.as_ptr();
+    let original_capacity = vec.capacity();
 
     let mut erased = TypeErasedVec::new(vec);
     assert_erased_state(&erased, 1, Layout::new::<DropTracker>());
@@ -469,7 +471,17 @@ fn test_guard_with_unwind_safety() {
 
     assert!(res.is_err());
 
-    assert_erased_state(&erased, 0, Layout::new::<DropTracker>());
+    assert_erased_state(&erased, 2, Layout::new::<DropTracker>());
+    assert_eq!(erased.capacity(), original_capacity);
+    assert_eq!(drop_count.get(), 0);
+
+    {
+        // SAFETY: The restoration guard retained the DropTracker element type.
+        let guard = unsafe { erased.cast_type::<DropTracker>() };
+        assert_eq!(guard.as_slice().as_ptr(), original_ptr);
+    }
+
+    drop(erased);
     assert_eq!(drop_count.get(), 2);
 }