command based entry api with EntityCommands::entry (#15274)

# Objective

It's convenient to be able to modify a component if it exist, and insert
a default value if it doesn't. You can already do most of this with
`EntityCommands::insert_if_new`, and all of this using a custom command.
However, that does not spark joy in my opinion.

Closes #10669

## Solution

Introduce a new commands type `EntityEntryCommands`, along with a method
to access it, `EntityCommands::entry`.

`EntityEntryCommands` exposes a subset of the entry API (`and_modify`,
`or_insert`, etc), however it's not an enum so it doesn't allow pattern
matching. Also, `or_insert` won't return the component because it's all
based on commands.

## Testing

Added a new test `entity_commands_entry`.

---

## Showcase

```rust
commands
    .entity(player)
    .entry::<Level>()
    .and_modify(|mut lvl| lvl.0 += 1)
    .or_default();
```

---------

Co-authored-by: Jan Hohenheim <jan@hohenheim.ch>
This commit is contained in:
poopy 2024-09-19 17:20:13 +02:00 committed by GitHub
parent 1bb8007dce
commit 3a66d88c83
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1,12 +1,14 @@
mod parallel_scope; mod parallel_scope;
use core::panic::Location; use core::panic::Location;
use std::marker::PhantomData;
use super::{Deferred, IntoObserverSystem, IntoSystem, RegisterSystem, Resource}; use super::{Deferred, IntoObserverSystem, IntoSystem, RegisterSystem, Resource};
use crate::{ use crate::{
self as bevy_ecs, self as bevy_ecs,
bundle::{Bundle, InsertMode}, bundle::{Bundle, InsertMode},
component::{ComponentId, ComponentInfo}, change_detection::Mut,
component::{Component, ComponentId, ComponentInfo},
entity::{Entities, Entity}, entity::{Entities, Entity},
event::{Event, SendEvent}, event::{Event, SendEvent},
observer::{Observer, TriggerEvent, TriggerTargets}, 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. /// Adds a [`Bundle`] of components to the entity.
/// ///
/// This will overwrite any previous value(s) of the same component type. /// 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 /// components will leave the old values instead of replacing them with new
/// ones. /// 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 /// # Panics
/// ///
/// The command will panic when applied if the associated entity does not exist. /// 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 impl<F> Command for F
where where
F: FnOnce(&mut World) + Send + 'static, 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. /// An [`EntityCommand`] that attempts to add the components in a [`Bundle`] to an entity.
/// Does nothing if the entity does not exist. /// Does nothing if the entity does not exist.
#[track_caller] #[track_caller]
@ -1660,7 +1823,7 @@ mod tests {
self as bevy_ecs, self as bevy_ecs,
component::Component, component::Component,
system::{Commands, Resource}, system::{Commands, Resource},
world::{CommandQueue, World}, world::{CommandQueue, FromWorld, World},
}; };
use std::{ use std::{
any::TypeId, any::TypeId,
@ -1697,6 +1860,50 @@ mod tests {
world.spawn((W(0u32), W(42u64))); 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] #[test]
fn commands() { fn commands() {
let mut world = World::default(); let mut world = World::default();