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,