fallback to remove components one by one when failing to remove a bundle

crates/bevy_ecs/hecs/src/world.rs

@@ -612,17 +612,35 @@ impl World {
     /// assert_eq!(*world.get::<bool>(e).unwrap(), true);
     /// ```
     pub fn remove<T: Bundle>(&mut self, entity: Entity) -> Result<T, ComponentError> {
+        self.flush();
+        let loc = self.entities.get_mut(entity)?;
+        unsafe {
+            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;
 
-        self.flush();
         let loc = self.entities.get_mut(entity)?;
         unsafe {
-            let removed = T::with_static_ids(|ids| ids.iter().copied().collect::<HashSet<_>>());
             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) {
@@ -636,8 +654,6 @@ 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,
@@ -664,10 +680,37 @@ impl World {
             {
                 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`
     ///
     /// See `remove`.

crates/bevy_ecs/src/system/commands.rs

@@ -126,7 +126,20 @@ where
     T: Bundle + Send + Sync + 'static,
 {
     fn write(self: Box<Self>, world: &mut World, _resources: &mut Resources) {
-        world.remove::<T>(self.entity).unwrap();
+        if let Err(e) = world.remove::<T>(self.entity) {
+            log::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) {
+                log::debug!(
+                    "Failed to remove components {:?} with error: {}",
+                    std::any::type_name::<T>(),
+                    e
+                );
+            }
+        }
     }
 }
 
@@ -373,4 +386,34 @@ 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)>()
+            .iter()
+            .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)>()
+            .iter()
+            .map(|(a, b)| (*a, *b))
+            .collect::<Vec<_>>();
+        assert_eq!(results_after, vec![]);
+        let results_after_u64 = world.query::<&u64>().iter().map(|a| *a).collect::<Vec<_>>();
+        assert_eq!(results_after_u64, vec![]);
+    }
 }