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

crates/bevy_ecs/hecs/src/world.rs

@@ -671,62 +671,102 @@ 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 info = self.archetypes[loc.archetype as usize]
-                .types()
-                .iter()
-                .cloned()
-                .filter(|x| !removed.contains(&x.id()))
-                .collect::<Vec<_>>();
-            let elements = info.iter().map(|x| x.id()).collect::<Vec<_>>();
-            let target = match self.index.entry(elements) {
-                Entry::Occupied(x) => *x.get(),
-                Entry::Vacant(x) => {
-                    self.archetypes.push(Archetype::new(info));
-                    let index = (self.archetypes.len() - 1) as u32;
-                    x.insert(index);
-                    self.archetype_generation += 1;
-                    index
-                }
-            };
+            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))?;
-            let (source_arch, target_arch) = index2(
-                &mut self.archetypes,
-                loc.archetype as usize,
-                target as usize,
-            );
-            let target_index = target_arch.allocate(entity);
-            loc.archetype = target;
-            loc.index = target_index;
-            let removed_components = &mut self.removed_components;
-            if let Some(moved) =
-                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) {
-                        ptr::copy_nonoverlapping(src, dst.as_ptr(), size);
-                        let state = target_arch.get_type_state_mut(ty).unwrap();
-                        *state.added().as_ptr().add(target_index) = is_added;
-                        *state.mutated().as_ptr().add(target_index) = is_mutated;
-                    } else {
-                        let removed_entities =
-                            removed_components.entry(ty).or_insert_with(Vec::new);
-                        removed_entities.push(entity);
-                    }
-                })
-            {
-                self.entities.get_mut(moved).unwrap().index = old_index;
+            match self.remove_bundle_internal(entity, to_remove) {
+                Ok(_) => Ok(bundle),
+                Err(err) => Err(err),
             }
-            Ok(bundle)
         }
     }
 
+    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| !to_remove.contains(&x.id()))
+            .collect::<Vec<_>>();
+        let elements = info.iter().map(|x| x.id()).collect::<Vec<_>>();
+        let target = match self.index.entry(elements) {
+            Entry::Occupied(x) => *x.get(),
+            Entry::Vacant(x) => {
+                self.archetypes.push(Archetype::new(info));
+                let index = (self.archetypes.len() - 1) as u32;
+                x.insert(index);
+                self.archetype_generation += 1;
+                index
+            }
+        };
+        let old_index = loc.index;
+        let (source_arch, target_arch) = index2(
+            &mut self.archetypes,
+            loc.archetype as usize,
+            target as usize,
+        );
+        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) = 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) {
+                    ptr::copy_nonoverlapping(src, dst.as_ptr(), size);
+                    let state = target_arch.get_type_state_mut(ty).unwrap();
+                    *state.added().as_ptr().add(target_index) = is_added;
+                    *state.mutated().as_ptr().add(target_index) = is_mutated;
+                } else {
+                    let removed_entities = removed_components.entry(ty).or_insert_with(Vec::new);
+                    removed_entities.push(entity);
+                }
+            })
+        } {
+            self.entities.get_mut(moved).unwrap().index = old_index;
+        }
+        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

@@ -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![]);
+    }
 }