Fall back to remove components one by one when failing to remove a bundle (#719)

Fall back to remove components one by one when failing to remove a bundle

crates/bevy_ecs/hecs/src/world.rs

@@ -671,17 +671,34 @@ impl World {
     /// assert_eq!(*world.get::<bool>(e).unwrap(), true);
     /// ```
     pub fn remove<T: Bundle>(&mut self, entity: Entity) -> Result<T, ComponentError> {
-        use std::collections::hash_map::Entry;
-
         self.flush();
         let loc = self.entities.get_mut(entity)?;
         unsafe {
-            let removed = T::with_static_ids(|ids| ids.iter().copied().collect::<HashSet<_>>());
+            let to_remove = T::with_static_ids(|ids| ids.iter().copied().collect::<HashSet<_>>());
+
+            let old_index = loc.index;
+            let source_arch = &self.archetypes[loc.archetype as usize];
+            let bundle = T::get(|ty, size| source_arch.get_dynamic(ty, size, old_index))?;
+            match self.remove_bundle_internal(entity, to_remove) {
+                Ok(_) => Ok(bundle),
+                Err(err) => Err(err),
+            }
+        }
+    }
+
+    fn remove_bundle_internal<S: core::hash::BuildHasher>(
+        &mut self,
+        entity: Entity,
+        to_remove: std::collections::HashSet<TypeId, S>,
+    ) -> Result<(), ComponentError> {
+        use std::collections::hash_map::Entry;
+
+        let loc = self.entities.get_mut(entity)?;
         let info = self.archetypes[loc.archetype as usize]
             .types()
             .iter()
             .cloned()
-                .filter(|x| !removed.contains(&x.id()))
+            .filter(|x| !to_remove.contains(&x.id()))
             .collect::<Vec<_>>();
         let elements = info.iter().map(|x| x.id()).collect::<Vec<_>>();
         let target = match self.index.entry(elements) {
@@ -695,18 +712,16 @@ impl World {
             }
         };
         let old_index = loc.index;
-            let source_arch = &self.archetypes[loc.archetype as usize];
-            let bundle = T::get(|ty, size| source_arch.get_dynamic(ty, size, old_index))?;
         let (source_arch, target_arch) = index2(
             &mut self.archetypes,
             loc.archetype as usize,
             target as usize,
         );
-            let target_index = target_arch.allocate(entity);
+        let target_index = unsafe { target_arch.allocate(entity) };
         loc.archetype = target;
         loc.index = target_index;
         let removed_components = &mut self.removed_components;
-            if let Some(moved) =
+        if let Some(moved) = unsafe {
             source_arch.move_to(old_index, |src, ty, size, is_added, is_mutated| {
                 // Only move the components present in the target archetype, i.e. the non-removed ones.
                 if let Some(dst) = target_arch.get_dynamic(ty, size, target_index) {
@@ -715,16 +730,41 @@ impl World {
                     *state.added().as_ptr().add(target_index) = is_added;
                     *state.mutated().as_ptr().add(target_index) = is_mutated;
                 } else {
-                        let removed_entities =
+                    let removed_entities = removed_components.entry(ty).or_insert_with(Vec::new);
-                            removed_components.entry(ty).or_insert_with(Vec::new);
                     removed_entities.push(entity);
                 }
             })
-            {
+        } {
             self.entities.get_mut(moved).unwrap().index = old_index;
         }
-            Ok(bundle)
+        Ok(())
     }
+
+    /// Remove components from `entity`
+    ///
+    /// Fallback method for `remove` when one of the component in `T` is not present in `entity`.
+    /// In this case, the missing component will be skipped.
+    ///
+    /// See `remove`.
+    pub fn remove_one_by_one<T: Bundle>(&mut self, entity: Entity) -> Result<(), ComponentError> {
+        self.flush();
+
+        let to_remove = T::with_static_ids(|ids| ids.iter().copied().collect::<HashSet<_>>());
+        for component_to_remove in to_remove.into_iter() {
+            let loc = self.entities.get(entity)?;
+            if loc.archetype == 0 {
+                return Err(ComponentError::NoSuchEntity);
+            }
+            if self.archetypes[loc.archetype as usize].has_dynamic(component_to_remove) {
+                let mut single_component_hashset = std::collections::HashSet::new();
+                single_component_hashset.insert(component_to_remove);
+                match self.remove_bundle_internal(entity, single_component_hashset) {
+                    Ok(_) | Err(ComponentError::MissingComponent(_)) => (),
+                    Err(err) => return Err(err),
+                };
+            }
+        }
+        Ok(())
     }
 
     /// Remove the `T` component from `entity`

crates/bevy_ecs/src/system/commands.rs

@@ -1,7 +1,7 @@
 use super::SystemId;
 use crate::resource::{Resource, Resources};
 use bevy_hecs::{Bundle, Component, DynamicBundle, Entity, EntityReserver, World};
-use bevy_utils::tracing::debug;
+use bevy_utils::tracing::{debug, warn};
 use std::marker::PhantomData;
 
 /// A [World] mutation
@@ -126,7 +126,30 @@ where
     T: Bundle + Send + Sync + 'static,
 {
     fn write(self: Box<Self>, world: &mut World, _resources: &mut Resources) {
-        world.remove::<T>(self.entity).unwrap();
+        match world.remove::<T>(self.entity) {
+            Ok(_) => (),
+            Err(bevy_hecs::ComponentError::MissingComponent(e)) => {
+                warn!(
+                    "Failed to remove components {:?} with error: {}. Falling back to inefficient one-by-one component removing.",
+                    std::any::type_name::<T>(),
+                    e
+                );
+                if let Err(e) = world.remove_one_by_one::<T>(self.entity) {
+                    debug!(
+                        "Failed to remove components {:?} with error: {}",
+                        std::any::type_name::<T>(),
+                        e
+                    );
+                }
+            }
+            Err(e) => {
+                debug!(
+                    "Failed to remove components {:?} with error: {}",
+                    std::any::type_name::<T>(),
+                    e
+                );
+            }
+        }
     }
 }
 
@@ -329,4 +352,32 @@ mod tests {
             .collect::<Vec<_>>();
         assert_eq!(results2, vec![]);
     }
+
+    #[test]
+    fn remove_components() {
+        let mut world = World::default();
+        let mut resources = Resources::default();
+        let mut command_buffer = Commands::default();
+        command_buffer.set_entity_reserver(world.get_entity_reserver());
+        command_buffer.spawn((1u32, 2u64));
+        let entity = command_buffer.current_entity().unwrap();
+        command_buffer.apply(&mut world, &mut resources);
+        let results_before = world
+            .query::<(&u32, &u64)>()
+            .map(|(a, b)| (*a, *b))
+            .collect::<Vec<_>>();
+        assert_eq!(results_before, vec![(1u32, 2u64)]);
+
+        // test component removal
+        command_buffer.remove_one::<u32>(entity);
+        command_buffer.remove::<(u32, u64)>(entity);
+        command_buffer.apply(&mut world, &mut resources);
+        let results_after = world
+            .query::<(&u32, &u64)>()
+            .map(|(a, b)| (*a, *b))
+            .collect::<Vec<_>>();
+        assert_eq!(results_after, vec![]);
+        let results_after_u64 = world.query::<&u64>().map(|a| *a).collect::<Vec<_>>();
+        assert_eq!(results_after_u64, vec![]);
+    }
 }