diff --git a/src/guard.rs b/src/guard.rs index 8f4a6ae..b477a77 100644 --- a/src/guard.rs +++ b/src/guard.rs @@ -12,6 +12,33 @@ pub struct ContentGuard<'vec, T> { _phantom: PhantomData, } +/// 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>, + reerase: fn(Vec) -> TypeErasedVec, +} + +impl VecRestoreGuard<'_, T> { + fn vec_mut(&mut self) -> &mut Vec { + &mut self.vec + } +} + +impl 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 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(&mut self, f: impl FnOnce(&mut Vec) -> R) -> R { - let mut vec = self.take(); - // This is unwind-safe. If the closure panics, the inner `Vec` 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. diff --git a/src/tests.rs b/src/tests.rs index 8757fbf..771c4fd 100644 --- a/src/tests.rs +++ b/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::()); @@ -469,7 +471,17 @@ fn test_guard_with_unwind_safety() { assert!(res.is_err()); - assert_erased_state(&erased, 0, Layout::new::()); + assert_erased_state(&erased, 2, Layout::new::()); + 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::() }; + assert_eq!(guard.as_slice().as_ptr(), original_ptr); + } + + drop(erased); assert_eq!(drop_count.get(), 2); }