From eb6afd26a168f0252e856eea096232af0ff834b2 Mon Sep 17 00:00:00 2001 From: JaySpruce Date: Mon, 7 Jul 2025 15:02:55 -0500 Subject: [PATCH] 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> --- crates/bevy_ecs/src/relationship/mod.rs | 2 +- .../src/system/commands/entity_command.rs | 30 ++++++++++++- crates/bevy_ecs/src/system/commands/mod.rs | 45 ++++++++++++++++--- 3 files changed, 69 insertions(+), 8 deletions(-) diff --git a/crates/bevy_ecs/src/relationship/mod.rs b/crates/bevy_ecs/src/relationship/mod.rs index f95214262b..82b39e04e5 100644 --- a/crates/bevy_ecs/src/relationship/mod.rs +++ b/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 diff --git a/crates/bevy_ecs/src/system/commands/entity_command.rs b/crates/bevy_ecs/src/system/commands/entity_command.rs index 098493a148..6d977e808d 100644 --- a/crates/bevy_ecs/src/system/commands/entity_command.rs +++ b/crates/bevy_ecs/src/system/commands/entity_command.rs @@ -143,12 +143,38 @@ pub unsafe fn insert_by_id( /// 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(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::()) { + 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(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::()) { + let value = component_fn(); + entity.insert_with_caller(value, mode, caller, RelationshipHookMode::Run); + } } } diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index d345368ab2..f38c7b9b14 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/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(&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(&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 { + 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::>().entity().id(); queue.apply(&mut world); assert_eq!(id, entity); + let mut commands = Commands::new(&mut queue, &world); + commands + .entity(entity) + .entry::>() + .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::>(entity).unwrap().0); } #[test]