[Merged by Bors] - Improve safety for BlobVec::replace_unchecked

crates/bevy_ecs/src/storage/blob_vec.rs

@@ -6,6 +6,7 @@ use std::{
 };
 
 use bevy_ptr::{OwningPtr, Ptr, PtrMut};
+use bevy_utils::OnDrop;
 
 /// A flat, type-erased data storage type
 ///
@@ -153,30 +154,49 @@ impl BlobVec {
     /// [`BlobVec`]'s `item_layout`
     pub unsafe fn replace_unchecked(&mut self, index: usize, value: OwningPtr<'_>) {
         debug_assert!(index < self.len());
-        // If `drop` panics, then when the collection is dropped during stack unwinding, the
-        // collection's `Drop` impl will call `drop` again for the old value (which is still stored
-        // in the collection), so we get a double drop. To prevent that, we set len to 0 until we're
-        // done.
+
+        let drop = self.drop;
+        let size = self.item_layout.size();
+
+        // Temporarily set the length to zero, so that uninitialized bytes won't
+        // get observed in case this function panics.
         let old_len = self.len;
-        let ptr = self.get_unchecked_mut(index).promote().as_ptr();
         self.len = 0;
-        // Drop the old value, then write back, justifying the promotion
-        // If the drop impl for the old value panics then we run the drop impl for `value` too.
-        if let Some(drop) = self.drop {
-            struct OnDrop<F: FnMut()>(F);
-            impl<F: FnMut()> Drop for OnDrop<F> {
-                fn drop(&mut self) {
-                    (self.0)();
-                }
-            }
-            let value = value.as_ptr();
-            let on_unwind = OnDrop(|| (drop)(OwningPtr::new(NonNull::new_unchecked(value))));
 
-            (drop)(OwningPtr::new(NonNull::new_unchecked(ptr)));
+        // Transfer ownership of the old value out of the vector, so it can be dropped.
+        // SAFETY:
+        // * The caller has promized that `index` fits in this vector.
+        // * The storage location will get overwritten with `value` later, which ensures
+        //   that the element will not get observed or double dropped later.
+        // * If a panic occurs, `self.len` will remain `0`, which ensures a double-drop
+        //   does not occur. Instead, all elements will be forgotten.
+        let old_value = self.get_ptr_mut().byte_add(index * size).promote();
+
+        let dst_addr = old_value.as_ptr();
+        let src_addr = value.as_ptr();
+
+        if let Some(drop) = drop {
+            // This closusure will run in case `drop()` panics, which is similar to
+            // the `try ... catch` construct in languages like C++.
+            // This ensures that `value` does not get forgotten in case of a panic.
+            let on_unwind = OnDrop::new(|| drop(value));
+
+            drop(old_value);
 
+            // If the above code does not panic, make sure that `value` doesn't get dropped.
             core::mem::forget(on_unwind);
         }
-        std::ptr::copy_nonoverlapping::<u8>(value.as_ptr(), ptr, self.item_layout.size());
+        // Copy the new value into the vector, overwriting the previous value which has been moved.
+        // SAFETY:
+        // - `src_addr` and `dst_addr` were obtained from `OwningPtr`s, which ensures they are
+        //   valid for both reads and writes.
+        // - The data behind `src_addr` will only be dropped if the above branch panics,
+        //   so it must still be initialized and it is safe to transfer ownership into the vector.
+        // - `src_addr` and `dst_addr` were obtained from different memory locations,
+        //   both of which we have exclusive access to, so they are guaranteed not to overlap.
+        std::ptr::copy_nonoverlapping::<u8>(src_addr, dst_addr, size);
+
+        // Now that each entry in the vector is fully initialized again, restore its length.
         self.len = old_len;
     }
 

crates/bevy_utils/src/lib.rs

@@ -34,6 +34,7 @@ use std::{
     future::Future,
     hash::{BuildHasher, Hash, Hasher},
     marker::PhantomData,
+    mem::ManuallyDrop,
     ops::Deref,
     pin::Pin,
 };
@@ -233,3 +234,26 @@ impl<K: Hash + Eq + PartialEq + Clone, V> PreHashMapExt<K, V> for PreHashMap<K,
         }
     }
 }
+
+/// A type which calls a function when dropped.
+/// This can be used to ensure that cleanup code is run even in case of a panic.
+pub struct OnDrop<F: FnOnce()> {
+    callback: ManuallyDrop<F>,
+}
+
+impl<F: FnOnce()> OnDrop<F> {
+    /// Returns an object that will invoke the specified callback when dropped.
+    pub fn new(callback: F) -> Self {
+        Self {
+            callback: ManuallyDrop::new(callback),
+        }
+    }
+}
+
+impl<F: FnOnce()> Drop for OnDrop<F> {
+    fn drop(&mut self) {
+        // SAFETY: We may move out of `self`, since this instance can never be observed after it's dropped.
+        let callback = unsafe { ManuallyDrop::take(&mut self.callback) };
+        callback();
+    }
+}