diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index 9b135a122f..9933243cf8 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/crates/bevy_ecs/src/system/commands/mod.rs @@ -1,12 +1,14 @@ mod parallel_scope; use core::panic::Location; +use std::marker::PhantomData; use super::{Deferred, IntoObserverSystem, IntoSystem, RegisterSystem, Resource}; use crate::{ self as bevy_ecs, bundle::{Bundle, InsertMode}, - component::{ComponentId, ComponentInfo}, + change_detection::Mut, + component::{Component, ComponentId, ComponentInfo}, entity::{Entities, Entity}, event::{Event, SendEvent}, observer::{Observer, TriggerEvent, TriggerTargets}, @@ -906,6 +908,38 @@ impl EntityCommands<'_> { } } + /// Get an [`EntityEntryCommands`] for the [`Component`] `T`, + /// allowing you to modify it or insert it if it isn't already present. + /// + /// See also [`insert_if_new`](Self::insert_if_new), which lets you insert a [`Bundle`] without overwriting it. + /// + /// # Example + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # #[derive(Resource)] + /// # struct PlayerEntity { entity: Entity } + /// #[derive(Component)] + /// struct Level(u32); + /// + /// fn level_up_system(mut commands: Commands, player: Res) { + /// commands + /// .entity(player.entity) + /// .entry::() + /// // Modify the component if it exists + /// .and_modify(|mut lvl| lvl.0 += 1) + /// // Otherwise insert a default value + /// .or_insert(Level(0)); + /// } + /// # bevy_ecs::system::assert_is_system(level_up_system); + /// ``` + pub fn entry(&mut self) -> EntityEntryCommands { + EntityEntryCommands { + entity_commands: self.reborrow(), + marker: PhantomData, + } + } + /// Adds a [`Bundle`] of components to the entity. /// /// This will overwrite any previous value(s) of the same component type. @@ -1010,6 +1044,9 @@ impl EntityCommands<'_> { /// components will leave the old values instead of replacing them with new /// ones. /// + /// See also [`entry`](Self::entry), which lets you modify a [`Component`] if it's present, + /// as well as initialize it with a default value. + /// /// # Panics /// /// The command will panic when applied if the associated entity does not exist. @@ -1417,6 +1454,113 @@ impl EntityCommands<'_> { } } +/// A wrapper around [`EntityCommands`] with convenience methods for working with a specified component type. +pub struct EntityEntryCommands<'a, T> { + entity_commands: EntityCommands<'a>, + marker: PhantomData, +} + +impl<'a, T: Component> EntityEntryCommands<'a, T> { + /// Modify the component `T` if it exists, using the the function `modify`. + pub fn and_modify(mut self, modify: impl FnOnce(Mut) + Send + Sync + 'static) -> Self { + self.entity_commands = self + .entity_commands + .queue(move |mut entity: EntityWorldMut| { + if let Some(value) = entity.get_mut() { + modify(value); + } + }); + self + } + + /// [Insert](EntityCommands::insert) `default` into this entity, if `T` is not already present. + /// + /// See also [`or_insert_with`](Self::or_insert_with). + /// + /// # Panics + /// + /// Panics if the entity does not exist. + /// See [`or_try_insert`](Self::or_try_insert) for a non-panicking version. + #[track_caller] + pub fn or_insert(mut self, default: T) -> Self { + self.entity_commands = self + .entity_commands + .queue(insert(default, InsertMode::Keep)); + self + } + + /// [Insert](EntityCommands::insert) `default` into this entity, if `T` is not already present. + /// + /// Unlike [`or_insert`](Self::or_insert), this will not panic if the entity does not exist. + /// + /// See also [`or_insert_with`](Self::or_insert_with). + #[track_caller] + pub fn or_try_insert(mut self, default: T) -> Self { + self.entity_commands = self + .entity_commands + .queue(try_insert(default, InsertMode::Keep)); + self + } + + /// [Insert](EntityCommands::insert) the value returned from `default` into this entity, if `T` is not already present. + /// + /// See also [`or_insert`](Self::or_insert) and [`or_try_insert`](Self::or_try_insert). + /// + /// # Panics + /// + /// Panics if the entity does not exist. + /// See [`or_try_insert_with`](Self::or_try_insert_with) for a non-panicking version. + #[track_caller] + pub fn or_insert_with(self, default: impl Fn() -> T) -> Self { + self.or_insert(default()) + } + + /// [Insert](EntityCommands::insert) the value returned from `default` into this entity, if `T` is not already present. + /// + /// Unlike [`or_insert_with`](Self::or_insert_with), this will not panic if the entity does not exist. + /// + /// See also [`or_insert`](Self::or_insert) and [`or_try_insert`](Self::or_try_insert). + #[track_caller] + pub fn or_try_insert_with(self, default: impl Fn() -> T) -> Self { + self.or_try_insert(default()) + } + + /// [Insert](EntityCommands::insert) `T::default` into this entity, if `T` is not already present. + /// + /// See also [`or_insert`](Self::or_insert) and [`or_from_world`](Self::or_from_world). + /// + /// # Panics + /// + /// Panics if the entity does not exist. + #[track_caller] + pub fn or_default(self) -> Self + where + T: Default, + { + #[allow(clippy::unwrap_or_default)] + // FIXME: use `expect` once stable + self.or_insert(T::default()) + } + + /// [Insert](EntityCommands::insert) `T::from_world` into this entity, if `T` is not already present. + /// + /// See also [`or_insert`](Self::or_insert) and [`or_default`](Self::or_default). + /// + /// # Panics + /// + /// Panics if the entity does not exist. + #[track_caller] + pub fn or_from_world(mut self) -> Self + where + T: FromWorld, + { + self.entity_commands = self + .entity_commands + .queue(insert_from_world::(InsertMode::Keep)); + self + } +} + impl Command for F where F: FnOnce(&mut World) + Send + 'static, @@ -1525,6 +1669,25 @@ fn insert(bundle: T, mode: InsertMode) -> impl EntityCommand { } } +/// An [`EntityCommand`] that adds the component using its `FromWorld` implementation. +#[track_caller] +fn insert_from_world(mode: InsertMode) -> impl EntityCommand { + let caller = Location::caller(); + move |entity: Entity, world: &mut World| { + let value = T::from_world(world); + if let Some(mut entity) = world.get_entity_mut(entity) { + entity.insert_with_caller( + value, + mode, + #[cfg(feature = "track_change_detection")] + caller, + ); + } else { + panic!("error[B0003]: {caller}: Could not insert a bundle (of type `{}`) for entity {:?} because it doesn't exist in this World. See: https://bevyengine.org/learn/errors/b0003", std::any::type_name::(), entity); + } + } +} + /// An [`EntityCommand`] that attempts to add the components in a [`Bundle`] to an entity. /// Does nothing if the entity does not exist. #[track_caller] @@ -1660,7 +1823,7 @@ mod tests { self as bevy_ecs, component::Component, system::{Commands, Resource}, - world::{CommandQueue, World}, + world::{CommandQueue, FromWorld, World}, }; use std::{ any::TypeId, @@ -1697,6 +1860,50 @@ mod tests { world.spawn((W(0u32), W(42u64))); } + impl FromWorld for W { + fn from_world(world: &mut World) -> Self { + let v = world.resource::>(); + Self("*".repeat(v.0)) + } + } + + #[test] + fn entity_commands_entry() { + let mut world = World::default(); + let mut queue = CommandQueue::default(); + let mut commands = Commands::new(&mut queue, &world); + let entity = commands.spawn_empty().id(); + commands + .entity(entity) + .entry::>() + .and_modify(|_| unreachable!()); + queue.apply(&mut world); + assert!(!world.entity(entity).contains::>()); + let mut commands = Commands::new(&mut queue, &world); + commands + .entity(entity) + .entry::>() + .or_insert(W(0)) + .and_modify(|mut val| { + val.0 = 21; + }); + queue.apply(&mut world); + assert_eq!(21, world.get::>(entity).unwrap().0); + let mut commands = Commands::new(&mut queue, &world); + commands + .entity(entity) + .entry::>() + .and_modify(|_| unreachable!()) + .or_insert(W(42)); + queue.apply(&mut world); + assert_eq!(42, world.get::>(entity).unwrap().0); + world.insert_resource(W(5_usize)); + let mut commands = Commands::new(&mut queue, &world); + commands.entity(entity).entry::>().or_from_world(); + queue.apply(&mut world); + assert_eq!("*****", &world.get::>(entity).unwrap().0); + } + #[test] fn commands() { let mut world = World::default();