diff --git a/crates/bevy_ecs/src/reflect/entity_commands.rs b/crates/bevy_ecs/src/reflect/entity_commands.rs new file mode 100644 index 0000000000000..da9374bbf8a06 --- /dev/null +++ b/crates/bevy_ecs/src/reflect/entity_commands.rs @@ -0,0 +1,452 @@ +use crate::prelude::Mut; +use crate::reflect::AppTypeRegistry; +use crate::system::{Command, EntityCommands, Resource}; +use crate::{entity::Entity, reflect::ReflectComponent, world::World}; +use bevy_reflect::{Reflect, TypeRegistry}; +use std::borrow::Cow; +use std::marker::PhantomData; + +/// An extension trait for [`EntityCommands`] for reflection related functions +pub trait ReflectCommandExt { + /// Adds the given boxed reflect component to the entity using the reflection data in + /// [`AppTypeRegistry`]. + /// + /// This will overwrite any previous component of the same type. + /// + /// # Panics + /// + /// - If the entity doesn't exist. + /// - If [`AppTypeRegistry`] does not have the reflection data for the given [`Component`](crate::component::Component). + /// - If the component data is invalid. See [`Reflect::apply`] for further details. + /// - If [`AppTypeRegistry`] is not present in the [`World`]. + /// + /// # Note + /// + /// Prefer to use the typed [`EntityCommands::insert`] if possible. Adding a reflected component + /// is much slower. + /// + /// # Example + /// + /// ```rust + /// // Note that you need to register the component type in the AppTypeRegistry prior to using + /// // reflection. You can use the helpers on the App with `app.register_type::()` + /// // or write to the TypeRegistry directly to register all your components + /// + /// # use bevy_ecs::prelude::*; + /// # use bevy_ecs::reflect::ReflectCommandExt; + /// # use bevy_reflect::{FromReflect, FromType, Reflect, TypeRegistry}; + /// // A resource that can hold any component that implements reflect as a boxed reflect component + /// #[derive(Resource)] + /// struct Prefab{ + /// component: Box, + /// } + /// #[derive(Component, Reflect, Default)] + /// #[reflect(Component)] + /// struct ComponentA(u32); + /// + /// #[derive(Component, Reflect, Default)] + /// #[reflect(Component)] + /// struct ComponentB(String); + /// + /// fn insert_reflect_component( + /// mut commands: Commands, + /// mut prefab: ResMut + /// ) { + /// // Create a set of new boxed reflect components to use + /// let boxed_reflect_component_a: Box = Box::new(ComponentA(916)); + /// let boxed_reflect_component_b: Box = Box::new(ComponentB("NineSixteen".to_string())); + /// + /// // You can overwrite the component in the resource with either ComponentA or ComponentB + /// prefab.component = boxed_reflect_component_a; + /// prefab.component = boxed_reflect_component_b; + /// + /// // No matter which component is in the resource and without knowing the exact type, you can + /// // use the insert_reflect entity command to insert that component into an entity. + /// commands + /// .spawn_empty() + /// .insert_reflect(prefab.component.clone_value()); + /// } + /// + /// ``` + fn insert_reflect(&mut self, component: Box) -> &mut Self; + + /// Same as [`insert_reflect`](ReflectCommandExt::insert_reflect), but using the `T` resource as type registry instead of + /// `AppTypeRegistry`. + /// + /// # Panics + /// + /// - If the given [`Resource`] is not present in the [`World`]. + /// + /// # Note + /// + /// - The given [`Resource`] is removed from the [`World`] before the command is applied. + fn insert_reflect_with_registry>( + &mut self, + component: Box, + ) -> &mut Self; + + /// Removes from the entity the component with the given type name registered in [`AppTypeRegistry`]. + /// + /// Does nothing if the entity does not have a component of the same type, if [`AppTypeRegistry`] + /// does not contain the reflection data for the given component, or if the entity does not exist. + /// + /// # Note + /// + /// Prefer to use the typed [`EntityCommands::remove`] if possible. Removing a reflected component + /// is much slower. + /// + /// # Example + /// + /// ```rust + /// // Note that you need to register the component type in the AppTypeRegistry prior to using + /// // reflection. You can use the helpers on the App with `app.register_type::()` + /// // or write to the TypeRegistry directly to register all your components + /// + /// # use bevy_ecs::prelude::*; + /// # use bevy_ecs::reflect::ReflectCommandExt; + /// # use bevy_reflect::{FromReflect, FromType, Reflect, TypeRegistry}; + /// + /// // A resource that can hold any component that implements reflect as a boxed reflect component + /// #[derive(Resource)] + /// struct Prefab{ + /// entity: Entity, + /// component: Box, + /// } + /// #[derive(Component, Reflect, Default)] + /// #[reflect(Component)] + /// struct ComponentA(u32); + /// #[derive(Component, Reflect, Default)] + /// #[reflect(Component)] + /// struct ComponentB(String); + /// + /// fn remove_reflect_component( + /// mut commands: Commands, + /// prefab: Res + /// ) { + /// // Prefab can hold any boxed reflect component. In this case either + /// // ComponentA or ComponentB. No matter which component is in the resource though, + /// // we can attempt to remove any component of that same type from an entity. + /// commands.entity(prefab.entity) + /// .remove_reflect(prefab.component.type_name().to_owned()); + /// } + /// + /// ``` + fn remove_reflect(&mut self, component_type_name: impl Into>) -> &mut Self; + /// Same as [`remove_reflect`](ReflectCommandExt::remove_reflect), but using the `T` resource as type registry instead of + /// `AppTypeRegistry`. + fn remove_reflect_with_registry>( + &mut self, + component_type_name: impl Into>, + ) -> &mut Self; +} + +impl<'w, 's, 'a> ReflectCommandExt for EntityCommands<'w, 's, 'a> { + fn insert_reflect(&mut self, component: Box) -> &mut Self { + self.commands.add(InsertReflect { + entity: self.entity, + component, + }); + self + } + + fn insert_reflect_with_registry>( + &mut self, + component: Box, + ) -> &mut Self { + self.commands.add(InsertReflectWithRegistry:: { + entity: self.entity, + _t: PhantomData, + component, + }); + self + } + + fn remove_reflect(&mut self, component_type_name: impl Into>) -> &mut Self { + self.commands.add(RemoveReflect { + entity: self.entity, + component_type_name: component_type_name.into(), + }); + self + } + + fn remove_reflect_with_registry>( + &mut self, + component_type_name: impl Into>, + ) -> &mut Self { + self.commands.add(RemoveReflectWithRegistry:: { + entity: self.entity, + _t: PhantomData, + component_type_name: component_type_name.into(), + }); + self + } +} + +/// Helper function to add a reflect component to a given entity +fn insert_reflect( + world: &mut World, + entity: Entity, + type_registry: &TypeRegistry, + component: Box, +) { + let type_info = component.type_name(); + let Some(mut entity) = world.get_entity_mut(entity) else { + panic!("error[B0003]: Could not insert a reflected component (of type {}) for entity {entity:?} because it doesn't exist in this World.", component.type_name()); + }; + let Some(type_registration) = type_registry.get_with_name(type_info) else { + panic!("Could not get type registration (for component type {}) because it doesn't exist in the TypeRegistry.", component.type_name()); + }; + let Some(reflect_component) = type_registration.data::() else { + panic!("Could not get ReflectComponent data (for component type {}) because it doesn't exist in this TypeRegistration.", component.type_name()); + }; + reflect_component.insert(&mut entity, &*component); +} + +/// A [`Command`] that adds the boxed reflect component to an entity using the data in +/// [`AppTypeRegistry`]. +/// +/// See [`ReflectCommandExt::insert_reflect`] for details. +pub struct InsertReflect { + /// The entity on which the component will be inserted. + pub entity: Entity, + /// The reflect [Component](crate::component::Component) that will be added to the entity. + pub component: Box, +} + +impl Command for InsertReflect { + fn apply(self, world: &mut World) { + let registry = world.get_resource::().unwrap().clone(); + insert_reflect(world, self.entity, ®istry.read(), self.component); + } +} + +/// A [`Command`] that adds the boxed reflect component to an entity using the data in the provided +/// [`Resource`] that implements [`AsRef`]. +/// +/// See [`ReflectCommandExt::insert_reflect_with_registry`] for details. +pub struct InsertReflectWithRegistry> { + /// The entity on which the component will be inserted. + pub entity: Entity, + pub _t: PhantomData, + /// The reflect [Component](crate::component::Component) that will be added to the entity. + pub component: Box, +} + +impl> Command for InsertReflectWithRegistry { + fn apply(self, world: &mut World) { + world.resource_scope(|world, registry: Mut| { + let registry: &TypeRegistry = registry.as_ref().as_ref(); + insert_reflect(world, self.entity, registry, self.component); + }); + } +} + +/// Helper function to remove a reflect component from a given entity +fn remove_reflect( + world: &mut World, + entity: Entity, + type_registry: &TypeRegistry, + component_type_name: Cow<'static, str>, +) { + let Some(mut entity) = world.get_entity_mut(entity) else { + return; + }; + let Some(type_registration) = type_registry.get_with_name(&component_type_name) else { + return; + }; + let Some(reflect_component) = type_registration.data::() else { + return; + }; + reflect_component.remove(&mut entity); +} + +/// A [`Command`] that removes the component of the same type as the given component type name from +/// the provided entity. +/// +/// See [`ReflectCommandExt::remove_reflect`] for details. +pub struct RemoveReflect { + /// The entity from which the component will be removed. + pub entity: Entity, + /// The [Component](crate::component::Component) type name that will be used to remove a component + /// of the same type from the entity. + pub component_type_name: Cow<'static, str>, +} + +impl Command for RemoveReflect { + fn apply(self, world: &mut World) { + let registry = world.get_resource::().unwrap().clone(); + remove_reflect( + world, + self.entity, + ®istry.read(), + self.component_type_name, + ); + } +} + +/// A [`Command`] that removes the component of the same type as the given component type name from +/// the provided entity using the provided [`Resource`] that implements [`AsRef`]. +/// +/// See [`ReflectCommandExt::remove_reflect_with_registry`] for details. +pub struct RemoveReflectWithRegistry> { + /// The entity from which the component will be removed. + pub entity: Entity, + pub _t: PhantomData, + /// The [Component](crate::component::Component) type name that will be used to remove a component + /// of the same type from the entity. + pub component_type_name: Cow<'static, str>, +} + +impl> Command for RemoveReflectWithRegistry { + fn apply(self, world: &mut World) { + world.resource_scope(|world, registry: Mut| { + let registry: &TypeRegistry = registry.as_ref().as_ref(); + remove_reflect(world, self.entity, registry, self.component_type_name); + }); + } +} + +#[cfg(test)] +mod tests { + use crate::prelude::{AppTypeRegistry, ReflectComponent}; + use crate::reflect::ReflectCommandExt; + use crate::system::{Commands, SystemState}; + use crate::{self as bevy_ecs, component::Component, world::World}; + use bevy_ecs_macros::Resource; + use bevy_reflect::{Reflect, TypeRegistry}; + + #[derive(Resource)] + struct TypeRegistryResource { + type_registry: TypeRegistry, + } + + impl AsRef for TypeRegistryResource { + fn as_ref(&self) -> &TypeRegistry { + &self.type_registry + } + } + + #[derive(Component, Reflect, Default, PartialEq, Eq, Debug)] + #[reflect(Component)] + struct ComponentA(u32); + + #[test] + fn insert_reflected() { + let mut world = World::new(); + + let type_registry = AppTypeRegistry::default(); + { + let mut registry = type_registry.write(); + registry.register::(); + registry.register_type_data::(); + } + world.insert_resource(type_registry); + + let mut system_state: SystemState = SystemState::new(&mut world); + let mut commands = system_state.get_mut(&mut world); + + let entity = commands.spawn_empty().id(); + + let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box; + + commands + .entity(entity) + .insert_reflect(boxed_reflect_component_a); + system_state.apply(&mut world); + + assert_eq!( + world.entity(entity).get::(), + Some(&ComponentA(916)) + ); + } + + #[test] + fn insert_reflected_with_registry() { + let mut world = World::new(); + + let mut type_registry = TypeRegistryResource { + type_registry: TypeRegistry::new(), + }; + + type_registry.type_registry.register::(); + type_registry + .type_registry + .register_type_data::(); + world.insert_resource(type_registry); + + let mut system_state: SystemState = SystemState::new(&mut world); + let mut commands = system_state.get_mut(&mut world); + + let entity = commands.spawn_empty().id(); + + let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box; + + commands + .entity(entity) + .insert_reflect_with_registry::(boxed_reflect_component_a); + system_state.apply(&mut world); + + assert_eq!( + world.entity(entity).get::(), + Some(&ComponentA(916)) + ); + } + + #[test] + fn remove_reflected() { + let mut world = World::new(); + + let type_registry = AppTypeRegistry::default(); + { + let mut registry = type_registry.write(); + registry.register::(); + registry.register_type_data::(); + } + world.insert_resource(type_registry); + + let mut system_state: SystemState = SystemState::new(&mut world); + let mut commands = system_state.get_mut(&mut world); + + let entity = commands.spawn(ComponentA(0)).id(); + + let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box; + + commands + .entity(entity) + .remove_reflect(boxed_reflect_component_a.type_name().to_owned()); + system_state.apply(&mut world); + + assert_eq!(world.entity(entity).get::(), None); + } + + #[test] + fn remove_reflected_with_registry() { + let mut world = World::new(); + + let mut type_registry = TypeRegistryResource { + type_registry: TypeRegistry::new(), + }; + + type_registry.type_registry.register::(); + type_registry + .type_registry + .register_type_data::(); + world.insert_resource(type_registry); + + let mut system_state: SystemState = SystemState::new(&mut world); + let mut commands = system_state.get_mut(&mut world); + + let entity = commands.spawn(ComponentA(0)).id(); + + let boxed_reflect_component_a = Box::new(ComponentA(916)) as Box; + + commands + .entity(entity) + .remove_reflect_with_registry::( + boxed_reflect_component_a.type_name().to_owned(), + ); + system_state.apply(&mut world); + + assert_eq!(world.entity(entity).get::(), None); + } +} diff --git a/crates/bevy_ecs/src/reflect/mod.rs b/crates/bevy_ecs/src/reflect/mod.rs index bcfb1770b770a..66cfeacde0948 100644 --- a/crates/bevy_ecs/src/reflect/mod.rs +++ b/crates/bevy_ecs/src/reflect/mod.rs @@ -7,10 +7,12 @@ use crate::{entity::Entity, system::Resource}; use bevy_reflect::{impl_reflect_value, ReflectDeserialize, ReflectSerialize, TypeRegistryArc}; mod component; +mod entity_commands; mod map_entities; mod resource; pub use component::{ReflectComponent, ReflectComponentFns}; +pub use entity_commands::ReflectCommandExt; pub use map_entities::ReflectMapEntities; pub use resource::{ReflectResource, ReflectResourceFns}; diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index fac21c7c0ffc2..59b51ad74562a 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/crates/bevy_ecs/src/system/commands/mod.rs @@ -642,8 +642,8 @@ impl Command for WithEntity { /// A list of commands that will be run to modify an [entity](crate::entity). pub struct EntityCommands<'w, 's, 'a> { - entity: Entity, - commands: &'a mut Commands<'w, 's>, + pub(crate) entity: Entity, + pub(crate) commands: &'a mut Commands<'w, 's>, } impl<'w, 's, 'a> EntityCommands<'w, 's, 'a> {