Fix UB in ResourceData and ComponentSparseSet

Author: james7132

crates/bevy_ecs/src/storage/resource.rs

@@ -5,7 +5,7 @@ use crate::{
 };
 use bevy_ptr::{OwningPtr, Ptr, UnsafeCellDeref};
 use bevy_utils::prelude::DebugName;
-use core::{cell::UnsafeCell, mem::ManuallyDrop, panic::Location};
+use core::{cell::UnsafeCell, panic::Location};
 
 #[cfg(feature = "std")]
 use std::thread::ThreadId;
@@ -16,7 +16,7 @@ use std::thread::ThreadId;
 ///
 /// [`World`]: crate::world::World
 pub struct ResourceData<const SEND: bool> {
-    data: ManuallyDrop<BlobArray>,
+    data: BlobArray,
     is_present: bool,
     added_ticks: UnsafeCell<Tick>,
     changed_ticks: UnsafeCell<Tick>,
@@ -50,7 +50,7 @@ impl<const SEND: bool> Drop for ResourceData<SEND> {
         // been dropped. The validate_access call above will check that the
         // data is dropped on the thread it was inserted from.
         unsafe {
-            ManuallyDrop::drop(&mut self.data);
+            self.data.drop(1, self.is_present().into());
         }
     }
 }
@@ -379,7 +379,7 @@ impl<const SEND: bool> Resources<SEND> {
                 )
             };
             ResourceData {
-                data: ManuallyDrop::new(data),
+                data,
                 is_present: false,
                 added_ticks: UnsafeCell::new(Tick::new(0)),
                 changed_ticks: UnsafeCell::new(Tick::new(0)),

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -3,8 +3,7 @@ use crate::{
     component::{CheckChangeTicks, ComponentId, ComponentInfo, ComponentTicks, Tick, TickCells},
     entity::{Entity, EntityRow},
     query::DebugCheckedUnwrap,
-    storage::VecExtensions,
-    storage::{AbortOnPanic, TableRow, ThinColumn},
+    storage::{AbortOnPanic, TableRow, ThinColumn, VecExtensions},
 };
 use alloc::{boxed::Box, vec::Vec};
 use bevy_ptr::{OwningPtr, Ptr};
@@ -155,7 +154,7 @@ impl ComponentSparseSet {
     /// Returns `true` if the sparse set contains no component values.
     #[inline]
     pub fn is_empty(&self) -> bool {
-        self.entities.len() == 0
+        self.entities.is_empty()
     }
 
     /// Inserts the `entity` key and component `value` pair into this sparse
@@ -330,12 +329,12 @@ impl ComponentSparseSet {
         self.sparse.remove(entity.row()).map(|dense_index| {
             #[cfg(debug_assertions)]
             assert_eq!(entity, self.entities[dense_index.index()]);
-            let len = self.entities.len();
-            if dense_index.index() >= len - 1 {
+            let last = self.entities.len() - 1;
+            if dense_index.index() >= last {
                 #[cfg(debug_assertions)]
-                assert_eq!(dense_index.index(), len - 1);
+                assert_eq!(dense_index.index(), last);
                 // SAFETY: TODO
-                unsafe { self.entities.set_len(dense_index.index()) };
+                unsafe { self.entities.set_len(last) };
                 // SAFETY: TODO
                 unsafe {
                     self.dense
@@ -361,7 +360,7 @@ impl ComponentSparseSet {
                 // SAFETY: dense_index was just removed from `sparse`, which ensures that it is valid
                 unsafe {
                     self.dense
-                        .swap_remove_and_forget_unchecked(len, dense_index)
+                        .swap_remove_and_forget_unchecked_nonoverlapping(last, dense_index)
                 }
             }
         })
@@ -376,14 +375,14 @@ impl ComponentSparseSet {
             .map(|dense_index| {
                 #[cfg(debug_assertions)]
                 assert_eq!(entity, self.entities[dense_index.index()]);
-                let len = self.entities.len();
-                if dense_index.index() >= len - 1 {
+                let last = self.entities.len() - 1;
+                if dense_index.index() >= last {
                     #[cfg(debug_assertions)]
-                    assert_eq!(dense_index.index(), len - 1);
+                    assert_eq!(dense_index.index(), last);
                     // SAFETY: TODO
-                    unsafe { self.entities.set_len(dense_index.index()) };
+                    unsafe { self.entities.set_len(last) };
                     // SAFETY: TODO
-                    unsafe { self.dense.drop_last_component(dense_index.index()) };
+                    unsafe { self.dense.drop_last_component(last) };
                 } else {
                     // SAFETY: TODO
                     unsafe {
@@ -403,8 +402,8 @@ impl ComponentSparseSet {
                     // SAFETY: TODO
                     unsafe {
                         self.dense
-                            .swap_remove_and_drop_unchecked_nonoverlapping(len, dense_index);
-                    };
+                            .swap_remove_and_drop_unchecked_nonoverlapping(last, dense_index);
+                    }
                 }
             })
             .is_some()
@@ -415,6 +414,17 @@ impl ComponentSparseSet {
     }
 }
 
+impl Drop for ComponentSparseSet {
+    fn drop(&mut self) {
+        let len = self.entities.len();
+        self.entities.clear();
+        // SAFETY: `cap` and `len` are correct. `dense` is never accessed again after this call.
+        unsafe {
+            self.dense.drop(self.entities.capacity(), len);
+        }
+    }
+}
+
 /// A data structure that blends dense and sparse storage
 ///
 /// `I` is the type of the indices, while `V` is the type of data stored in the dense storage.

crates/bevy_ecs/src/storage/table/mod.rs

@@ -3,7 +3,7 @@ use crate::{
     component::{CheckChangeTicks, ComponentId, ComponentInfo, ComponentTicks, Components, Tick},
     entity::Entity,
     query::DebugCheckedUnwrap,
-    storage::{ImmutableSparseSet, SparseSet},
+    storage::{AbortOnPanic, ImmutableSparseSet, SparseSet},
 };
 use alloc::{boxed::Box, vec, vec::Vec};
 use bevy_platform::collections::HashMap;
@@ -182,15 +182,6 @@ pub struct Table {
     entities: Vec<Entity>,
 }
 
-struct AbortOnPanic;
-
-impl Drop for AbortOnPanic {
-    fn drop(&mut self) {
-        // Panicking while unwinding will force an abort.
-        panic!("Aborting due to allocator error");
-    }
-}
-
 impl Table {
     /// Fetches a read-only slice of the entities stored within the [`Table`].
     #[inline]
@@ -821,7 +812,7 @@ impl Drop for Table {
         let cap = self.capacity();
         self.entities.clear();
         for col in self.columns.values_mut() {
-            // SAFETY: `cap` and `len` are correct
+            // SAFETY: `cap` and `len` are correct. `col` is never accessed again after this call.
             unsafe {
                 col.drop(cap, len);
             }