Avoid early function invocation in EntityEntryCommands (#19978)

## Objective

Fixes #19884.

## Solution

- Add an internal entity command `insert_with`, which takes a function
returning a component and checks if the component would actually be
inserted before invoking the function.
- Add the same check to `insert_from_world`, since it's a similar
situation.
- Update the `or_insert_with`, `or_try_insert_with`, and `or_default`
methods on `EntityEntryCommands` to use the new command.

Since the function/closure returning the component now needs to be sent
into the command (rather than being invoked before the command is
created), the function now has `Send + 'static` bounds. Pretty typical
for command stuff, but I don't know how/if it'll affect existing users.

---------

Co-authored-by: Chris Russell <8494645+chescock@users.noreply.github.com>

Author: JaySpruce

crates/bevy_ecs/src/relationship/mod.rs

@@ -133,7 +133,7 @@ pub trait Relationship: Component + Sized {
                 .and_modify(move |mut relationship_target| {
                     relationship_target.collection_mut_risky().add(entity);
                 })
-                .or_insert_with(|| {
+                .or_insert_with(move || {
                     let mut target = Self::RelationshipTarget::with_capacity(1);
                     target.collection_mut_risky().add(entity);
                     target

crates/bevy_ecs/src/system/commands/entity_command.rs

@@ -143,12 +143,38 @@ pub unsafe fn insert_by_id<T: Send + 'static>(
 
 /// An [`EntityCommand`] that adds a component to an entity using
 /// the component's [`FromWorld`] implementation.
+///
+/// `T::from_world` will only be invoked if the component will actually be inserted.
+/// In other words, `T::from_world` will *not* be invoked if `mode` is [`InsertMode::Keep`]
+/// and the entity already has the component.
 #[track_caller]
 pub fn insert_from_world<T: Component + FromWorld>(mode: InsertMode) -> impl EntityCommand {
     let caller = MaybeLocation::caller();
     move |mut entity: EntityWorldMut| {
-        let value = entity.world_scope(|world| T::from_world(world));
-        entity.insert_with_caller(value, mode, caller, RelationshipHookMode::Run);
+        if !(mode == InsertMode::Keep && entity.contains::<T>()) {
+            let value = entity.world_scope(|world| T::from_world(world));
+            entity.insert_with_caller(value, mode, caller, RelationshipHookMode::Run);
+        }
+    }
+}
+
+/// An [`EntityCommand`] that adds a component to an entity using
+/// some function that returns the component.
+///
+/// The function will only be invoked if the component will actually be inserted.
+/// In other words, the function will *not* be invoked if `mode` is [`InsertMode::Keep`]
+/// and the entity already has the component.
+#[track_caller]
+pub fn insert_with<T: Component, F>(component_fn: F, mode: InsertMode) -> impl EntityCommand
+where
+    F: FnOnce() -> T + Send + 'static,
+{
+    let caller = MaybeLocation::caller();
+    move |mut entity: EntityWorldMut| {
+        if !(mode == InsertMode::Keep && entity.contains::<T>()) {
+            let value = component_fn();
+            entity.insert_with_caller(value, mode, caller, RelationshipHookMode::Run);
+        }
     }
 }
 

crates/bevy_ecs/src/system/commands/mod.rs

@@ -2273,35 +2273,53 @@ impl<'a, T: Component> EntityEntryCommands<'a, T> {
 
     /// [Insert](EntityCommands::insert) the value returned from `default` into this entity,
     /// if `T` is not already present.
+    ///
+    /// `default` will only be invoked if the component will actually be inserted.
     #[track_caller]
-    pub fn or_insert_with(&mut self, default: impl Fn() -> T) -> &mut Self {
-        self.or_insert(default())
+    pub fn or_insert_with<F>(&mut self, default: F) -> &mut Self
+    where
+        F: FnOnce() -> T + Send + 'static,
+    {
+        self.entity_commands
+            .queue(entity_command::insert_with(default, InsertMode::Keep));
+        self
     }
 
     /// [Insert](EntityCommands::insert) the value returned from `default` into this entity,
     /// if `T` is not already present.
     ///
+    /// `default` will only be invoked if the component will actually be inserted.
+    ///
     /// # Note
     ///
     /// If the entity does not exist when this command is executed,
     /// the resulting error will be ignored.
     #[track_caller]
-    pub fn or_try_insert_with(&mut self, default: impl Fn() -> T) -> &mut Self {
-        self.or_try_insert(default())
+    pub fn or_try_insert_with<F>(&mut self, default: F) -> &mut Self
+    where
+        F: FnOnce() -> T + Send + 'static,
+    {
+        self.entity_commands
+            .queue_silenced(entity_command::insert_with(default, InsertMode::Keep));
+        self
     }
 
     /// [Insert](EntityCommands::insert) `T::default` into this entity,
     /// if `T` is not already present.
+    ///
+    /// `T::default` will only be invoked if the component will actually be inserted.
     #[track_caller]
     pub fn or_default(&mut self) -> &mut Self
     where
         T: Default,
     {
-        self.or_insert(T::default())
+        self.or_insert_with(T::default)
     }
 
     /// [Insert](EntityCommands::insert) `T::from_world` into this entity,
     /// if `T` is not already present.
+    ///
+    /// `T::from_world` will only be invoked if the component will actually be inserted.
     #[track_caller]
     pub fn or_from_world(&mut self) -> &mut Self
     where
@@ -2396,6 +2414,12 @@ mod tests {
         }
     }
 
+    impl Default for W<u8> {
+        fn default() -> Self {
+            unreachable!()
+        }
+    }
+
     #[test]
     fn entity_commands_entry() {
         let mut world = World::default();
@@ -2435,6 +2459,17 @@ mod tests {
         let id = commands.entity(entity).entry::<W<u64>>().entity().id();
         queue.apply(&mut world);
         assert_eq!(id, entity);
+        let mut commands = Commands::new(&mut queue, &world);
+        commands
+            .entity(entity)
+            .entry::<W<u8>>()
+            .or_insert_with(|| W(5))
+            .or_insert_with(|| unreachable!())
+            .or_try_insert_with(|| unreachable!())
+            .or_default()
+            .or_from_world();
+        queue.apply(&mut world);
+        assert_eq!(5, world.get::<W<u8>>(entity).unwrap().0);
     }
 
     #[test]