diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index 7fede918fdf19d..9f1c46a2e7b2b3 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/crates/bevy_ecs/src/system/commands/mod.rs @@ -528,6 +528,85 @@ impl<'w, 's> Commands<'w, 's> { } } +/// A [`Command`] which gets executed for a given [`Entity`]. +/// +/// # Examples +/// +/// ``` +/// # use std::collections::HashSet; +/// # use bevy_ecs::prelude::*; +/// use bevy_ecs::system::EntityCommand; +/// # +/// # #[derive(Component, PartialEq)] +/// # struct Name(String); +/// # impl Name { +/// # fn new(s: String) -> Self { Name(s) } +/// # fn as_str(&self) -> &str { &self.0 } +/// # } +/// +/// #[derive(Resource, Default)] +/// struct Counter(i64); +/// +/// /// A `Command` which names an entity based on a global counter. +/// struct CountName; +/// +/// impl EntityCommand for CountName { +/// fn write(self, id: Entity, world: &mut World) { +/// // Get the current value of the counter, and increment it for next time. +/// let mut counter = world.resource_mut::(); +/// let i = counter.0; +/// counter.0 += 1; +/// +/// // Name the entity after the value of the counter. +/// world.entity_mut(id).insert(Name::new(format!("Entity #{i}"))); +/// } +/// } +/// +/// // App creation boilerplate omitted... +/// # let mut world = World::new(); +/// # world.init_resource::(); +/// # +/// # let mut setup_stage = SystemStage::single_threaded().with_system(setup); +/// # let mut assert_stage = SystemStage::single_threaded().with_system(assert_names); +/// # +/// # setup_stage.run(&mut world); +/// # assert_stage.run(&mut world); +/// +/// fn setup(mut commands: Commands) { +/// commands.spawn_empty().add(CountName); +/// commands.spawn_empty().add(CountName); +/// } +/// +/// fn assert_names(named: Query<&Name>) { +/// // We use a HashSet because we do not care about the order. +/// let names: HashSet<_> = named.iter().map(Name::as_str).collect(); +/// assert_eq!(names, HashSet::from_iter(["Entity #0", "Entity #1"])); +/// } +/// ``` +pub trait EntityCommand: Send + 'static { + fn write(self, id: Entity, world: &mut World); + /// Returns a [`Command`] which executes this [`EntityCommand`] for the given [`Entity`]. + fn with_entity(self, id: Entity) -> WithEntity + where + Self: Sized, + { + WithEntity { cmd: self, id } + } +} + +/// Turns an [`EntityCommand`] type into a [`Command`] type. +pub struct WithEntity { + cmd: C, + id: Entity, +} + +impl Command for WithEntity { + #[inline] + fn write(self, world: &mut World) { + self.cmd.write(self.id, world); + } +} + /// A list of commands that will be run to modify an [entity](crate::entity). pub struct EntityCommands<'w, 's, 'a> { entity: Entity, @@ -690,6 +769,27 @@ impl<'w, 's, 'a> EntityCommands<'w, 's, 'a> { }); } + /// Pushes an [`EntityCommand`] to the queue, which will get executed for the current [`Entity`]. + /// + /// # Examples + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # fn my_system(mut commands: Commands) { + /// commands + /// .spawn_empty() + /// // Closures with this signature implement `EntityCommand`. + /// .add(|id: Entity, world: &mut World| { + /// println!("Executed an EntityCommand for {id:?}"); + /// }); + /// # } + /// # bevy_ecs::system::assert_is_system(my_system); + /// ``` + pub fn add(&mut self, command: C) -> &mut Self { + self.commands.add(command.with_entity(self.entity)); + self + } + /// Logs the components of the entity at the info level. /// /// # Panics @@ -716,6 +816,15 @@ where } } +impl EntityCommand for F +where + F: FnOnce(Entity, &mut World) + Send + 'static, +{ + fn write(self, id: Entity, world: &mut World) { + self(id, world); + } +} + #[derive(Debug)] pub struct Spawn { pub bundle: T,