From 8221c7fa65edec884ea454f701ec272cbddc119c Mon Sep 17 00:00:00 2001
From: sapir <yasapir@gmail.com>
Date: Mon, 20 Sep 2021 03:25:06 +0300
Subject: [PATCH] Fix double drop in BlobVec::replace_unchecked (#2597)

---
 crates/bevy_ecs/src/storage/blob_vec.rs   | 14 +++++++++++---
 crates/bevy_ecs/src/storage/sparse_set.rs | 15 ++++++++++-----
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_ecs/src/storage/blob_vec.rs b/crates/bevy_ecs/src/storage/blob_vec.rs
index 49c92d92d1eb90..7ea3b2bc1e7f88 100644
--- a/crates/bevy_ecs/src/storage/blob_vec.rs
+++ b/crates/bevy_ecs/src/storage/blob_vec.rs
@@ -88,6 +88,7 @@ impl BlobVec {
     /// # Safety
     /// - index must be in bounds
     /// - memory must be reserved and uninitialized
+    /// - the contents of `*value` must not be dropped by the caller
     #[inline]
     pub unsafe fn initialize_unchecked(&mut self, index: usize, value: *mut u8) {
         debug_assert!(index < self.len());
@@ -97,12 +98,19 @@ impl BlobVec {
 
     /// # Safety
     /// - index must be in-bounds
-    //  - memory must be previously initialized
+    /// - memory must be previously initialized
+    /// - the contents of `*value` must not be dropped by the caller. (In fact,
+    ///   this function may overwrite them.)
     pub unsafe fn replace_unchecked(&mut self, index: usize, value: *mut u8) {
         debug_assert!(index < self.len());
         let ptr = self.get_unchecked(index);
-        (self.drop)(ptr);
-        std::ptr::copy_nonoverlapping(value, ptr, self.item_layout.size());
+        // We need to make sure that we don't still contain the old value before
+        // calling `drop`. Otherwise, if `drop` panics, then when the collection
+        // is dropped stack unwinding, its `Drop` impl will call `drop` again
+        // for the old value, so we get a double drop. Also, the new value
+        // wouldn't be dropped.
+        std::ptr::swap_nonoverlapping(value, ptr, self.item_layout.size());
+        (self.drop)(value);
     }
 
     /// increases the length by one (and grows the vec if needed) with uninitialized memory and
diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs
index 42357ec87c6a3e..8ad76046e6cdc8 100644
--- a/crates/bevy_ecs/src/storage/sparse_set.rs
+++ b/crates/bevy_ecs/src/storage/sparse_set.rs
@@ -123,13 +123,18 @@ impl ComponentSparseSet {
         self.dense.len() == 0
     }
 
-    /// Inserts the `entity` key and component `value` pair into this sparse set.
-    /// The caller is responsible for ensuring the value is not dropped. This collection will drop
-    /// the value when needed.
+    /// Inserts the `entity` key and component `value` pair into this sparse
+    /// set. This collection takes ownership of the contents of `value`, and
+    /// will drop the value when needed. Also, it may overwrite the contents of
+    /// the `value` pointer if convenient. The caller is responsible for
+    /// ensuring it does not drop `*value` after calling `insert`.
     ///
     /// # Safety
-    /// The `value` pointer must point to a valid address that matches the `Layout`
-    ///  inside the `ComponentInfo` given when constructing this sparse set.
+    /// * The `value` pointer must point to a valid address that matches the
+    ///   `Layout` inside the `ComponentInfo` given when constructing this
+    ///   sparse set.
+    /// * The caller is responsible for ensuring it does not drop `*value` after
+    ///   calling `insert`.
     pub unsafe fn insert(&mut self, entity: Entity, value: *mut u8, change_tick: u32) {
         if let Some(&dense_index) = self.sparse.get(entity) {
             self.dense.replace_unchecked(dense_index, value);