drop sparse components in remove_bundle command

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -180,10 +180,7 @@ impl ComponentSparseSet {
     /// returned).
     pub fn remove_and_forget(&mut self, entity: Entity) -> Option<*mut u8> {
         self.sparse.remove(entity).map(|dense_index| {
-            // SAFE: unique access to ticks
-            unsafe {
-                (*self.ticks.get()).swap_remove(dense_index);
-            }
+            self.ticks.get_mut().swap_remove(dense_index);
             self.entities.swap_remove(dense_index);
             let is_last = dense_index == self.dense.len() - 1;
             // SAFE: dense_index was just removed from `sparse`, which ensures that it is valid

crates/bevy_ecs/src/system/commands.rs

@@ -417,9 +417,29 @@ impl<T: Component> Command for RemoveResource<T> {
 #[allow(clippy::float_cmp, clippy::approx_constant)]
 mod tests {
     use crate::{
+        component::{ComponentDescriptor, StorageType},
         system::{CommandQueue, Commands},
         world::World,
     };
+    use std::sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    };
+
+    #[derive(Clone, Debug)]
+    struct DropCk(Arc<AtomicUsize>);
+    impl DropCk {
+        fn new_pair() -> (Self, Arc<AtomicUsize>) {
+            let atomic = Arc::new(AtomicUsize::new(0));
+            (DropCk(atomic.clone()), atomic)
+        }
+    }
+
+    impl Drop for DropCk {
+        fn drop(&mut self) {
+            self.0.as_ref().fetch_add(1, Ordering::Relaxed);
+        }
+    }
 
     #[test]
     fn commands() {
@@ -454,10 +474,20 @@ mod tests {
     #[test]
     fn remove_components() {
         let mut world = World::default();
+
+        struct DenseDropCk(DropCk);
+        world
+            .register_component(ComponentDescriptor::new::<DropCk>(StorageType::SparseSet))
+            .unwrap();
+
         let mut command_queue = CommandQueue::default();
+        let (dense_dropck, dense_is_dropped) = DropCk::new_pair();
+        let dense_dropck = DenseDropCk(dense_dropck);
+        let (sparse_dropck, sparse_is_dropped) = DropCk::new_pair();
+
         let entity = Commands::new(&mut command_queue, &world)
             .spawn()
-            .insert_bundle((1u32, 2u64))
+            .insert_bundle((1u32, 2u64, dense_dropck, sparse_dropck))
             .id();
         command_queue.apply(&mut world);
         let results_before = world
@@ -471,8 +501,14 @@ mod tests {
         Commands::new(&mut command_queue, &world)
             .entity(entity)
             .remove::<u32>()
-            .remove_bundle::<(u32, u64)>();
+            .remove_bundle::<(u32, u64, DenseDropCk, DropCk)>();
+
+        assert_eq!(dense_is_dropped.load(Ordering::Relaxed), 0);
+        assert_eq!(sparse_is_dropped.load(Ordering::Relaxed), 0);
         command_queue.apply(&mut world);
+        assert_eq!(dense_is_dropped.load(Ordering::Relaxed), 1);
+        assert_eq!(sparse_is_dropped.load(Ordering::Relaxed), 1);
+
         let results_after = world
             .query::<(&u32, &u64)>()
             .iter(&world)

crates/bevy_ecs/src/world/entity_ref.rs

@@ -325,7 +325,7 @@ impl<'w> EntityMut<'w> {
             T::from_components(|| {
                 let component_id = bundle_components.next().unwrap();
                 // SAFE: entity location is valid and table row is removed below
-                remove_component(
+                take_component(
                     components,
                     storages,
                     old_archetype,
@@ -406,17 +406,18 @@ impl<'w> EntityMut<'w> {
         let entity = self.entity;
         for component_id in bundle_info.component_ids.iter().cloned() {
             if old_archetype.contains(component_id) {
-                // SAFE: entity location is valid and table row is removed below
-                unsafe {
-                    remove_component(
-                        components,
-                        storages,
-                        old_archetype,
-                        removed_components,
-                        component_id,
-                        entity,
-                        old_location,
-                    );
+                removed_components
+                    .get_or_insert_with(component_id, Vec::new)
+                    .push(entity);
+
+                // Make sure to drop components stored in sparse sets.
+                // Dense components are dropped later in `move_to_and_drop_missing_unchecked`.
+                if let Some(StorageType::SparseSet) = old_archetype.get_storage_type(component_id) {
+                    storages
+                        .sparse_sets
+                        .get_mut(component_id)
+                        .unwrap()
+                        .remove(entity);
                 }
             }
         }
@@ -586,13 +587,18 @@ unsafe fn get_component_and_ticks(
     }
 }
 
+/// Moves component data out of storage.
+///
+/// This function leaves the underlying memory unchanged, but the component behind
+/// returned pointer is semantically owned by the caller and will not be dropped in its original location.
+///
 /// # Safety
-// `entity_location` must be within bounds of the given archetype and `entity` must exist inside the
-// archetype
-/// The relevant table row must be removed separately
-/// `component_id` must be valid
+/// - `entity_location` must be within bounds of the given archetype and `entity` must exist inside the archetype
+/// - `component_id` must be valid
+/// - The relevant table row **must be removed** by the caller once all components are taken
+/// - caller is responsible to drop component data behind returned pointer
 #[inline]
-unsafe fn remove_component(
+unsafe fn take_component(
     components: &Components,
     storages: &mut Storages,
     archetype: &Archetype,