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>
This commit is contained in:
JaySpruce 2025-07-07 15:02:55 -05:00 committed by GitHub
parent b0df3fb4d0
commit eb6afd26a1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 69 additions and 8 deletions

View File

@ -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

View File

@ -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);
}
}
}

View File

@ -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]