command based entry api with EntityCommands::entry

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<PlayerEntity>) {
+    ///     commands
+    ///         .entity(player.entity)
+    ///         .entry::<Level>()
+    ///         // 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<T: Component>(&mut self) -> EntityEntryCommands<T> {
+        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<T>,
+}
+
+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<T>) + 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::<T>(InsertMode::Keep));
+        self
+    }
+}
+
 impl<F> Command for F
 where
     F: FnOnce(&mut World) + Send + 'static,
@@ -1525,6 +1669,25 @@ fn insert<T: Bundle>(bundle: T, mode: InsertMode) -> impl EntityCommand {
     }
 }
 
+/// An [`EntityCommand`] that adds the component using its `FromWorld` implementation.
+#[track_caller]
+fn insert_from_world<T: Component + FromWorld>(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::<T>(), 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<String> {
+        fn from_world(world: &mut World) -> Self {
+            let v = world.resource::<W<usize>>();
+            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::<W<u32>>()
+            .and_modify(|_| unreachable!());
+        queue.apply(&mut world);
+        assert!(!world.entity(entity).contains::<W<u32>>());
+        let mut commands = Commands::new(&mut queue, &world);
+        commands
+            .entity(entity)
+            .entry::<W<u32>>()
+            .or_insert(W(0))
+            .and_modify(|mut val| {
+                val.0 = 21;
+            });
+        queue.apply(&mut world);
+        assert_eq!(21, world.get::<W<u32>>(entity).unwrap().0);
+        let mut commands = Commands::new(&mut queue, &world);
+        commands
+            .entity(entity)
+            .entry::<W<u64>>()
+            .and_modify(|_| unreachable!())
+            .or_insert(W(42));
+        queue.apply(&mut world);
+        assert_eq!(42, world.get::<W<u64>>(entity).unwrap().0);
+        world.insert_resource(W(5_usize));
+        let mut commands = Commands::new(&mut queue, &world);
+        commands.entity(entity).entry::<W<String>>().or_from_world();
+        queue.apply(&mut world);
+        assert_eq!("*****", &world.get::<W<String>>(entity).unwrap().0);
+    }
+
     #[test]
     fn commands() {
         let mut world = World::default();