diff --git a/crates/bevy_ecs/hecs/src/world.rs b/crates/bevy_ecs/hecs/src/world.rs index 4613ffba79baca..d55aa0dd5b9e45 100644 --- a/crates/bevy_ecs/hecs/src/world.rs +++ b/crates/bevy_ecs/hecs/src/world.rs @@ -612,17 +612,35 @@ impl World { /// assert_eq!(*world.get::(e).unwrap(), true); /// ``` pub fn remove(&mut self, entity: Entity) -> Result { + self.flush(); + let loc = self.entities.get_mut(entity)?; + unsafe { + let to_remove = T::with_static_ids(|ids| ids.iter().copied().collect::>()); + + 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( + &mut self, + entity: Entity, + to_remove: std::collections::HashSet, + ) -> 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::>()); 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::>(); let elements = info.iter().map(|x| x.id()).collect::>(); 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(&mut self, entity: Entity) -> Result<(), ComponentError> { + self.flush(); + + let to_remove = T::with_static_ids(|ids| ids.iter().copied().collect::>()); + 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`. diff --git a/crates/bevy_ecs/src/system/commands.rs b/crates/bevy_ecs/src/system/commands.rs index f0b2218009ec24..5631e0c1e758dc 100644 --- a/crates/bevy_ecs/src/system/commands.rs +++ b/crates/bevy_ecs/src/system/commands.rs @@ -126,7 +126,20 @@ where T: Bundle + Send + Sync + 'static, { fn write(self: Box, world: &mut World, _resources: &mut Resources) { - world.remove::(self.entity).unwrap(); + if let Err(e) = world.remove::(self.entity) { + log::warn!( + "Failed to remove components {:?} with error: {}. Falling back to inefficient one-by-one component removing.", + std::any::type_name::(), + e + ); + if let Err(e) = world.remove_one_by_one::(self.entity) { + log::debug!( + "Failed to remove components {:?} with error: {}", + std::any::type_name::(), + e + ); + } + } } } @@ -373,4 +386,34 @@ mod tests { .collect::>(); 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::>(); + assert_eq!(results_before, vec![(1u32, 2u64)]); + + // test component removal + command_buffer.remove_one::(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::>(); + assert_eq!(results_after, vec![]); + let results_after_u64 = world.query::<&u64>().iter().map(|a| *a).collect::>(); + assert_eq!(results_after_u64, vec![]); + } }