Add an Entry api to EntityWorldMut. (#10650)

# Objective

Adds `.entry` to `EntityWorldMut` with `Entry`, `OccupiedEntry` and
`VacantEntry` for easier in-situ modification, based on `HashMap.entry`.

Fixes #10635

## Solution

This adds the `entry` method to `EntityWorldMut` which returns an
`Entry`. This is an enum of `OccupiedEntry` and `VacantEntry` and has
the methods `and_modify`, `insert_entry`, `or_insert`, `or_insert_with`
and `or_default`. The only difference between `OccupiedEntry` and
`VacantEntry` is the type, they are both a mutable reference to the
`EntityWorldMut` and a marker for the component type, `HashMap` also
stores things to make it quicker to access the data in `OccupiedEntry`
but I wasn't sure if we had anything it would be logical to store to
make accessing/modifying the component faster? As such, the differences
are that `OccupiedEntry` assumes the entity has the component (because
nothing else can have an `EntityWorldMut` so it can't be changed outside
the entry api) and has different methods.

All the methods are based very closely off `hashbrown::HashMap` (because
its easier to read the source of) with a couple of quirks like
`OccupiedEntry.insert` doesn't return the old value because we don't
appear to have an api for mem::replacing components.

---

## Changelog

- Added a new function `EntityWorldMut.entry` which returns an `Entry`,
allowing easier in-situ modification of a component.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Pascal Hertleif <killercup@gmail.com>
This commit is contained in:
TheBigCheese 2023-11-21 01:26:06 +00:00 committed by GitHub
parent 67df504671
commit 865041de74
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 387 additions and 2 deletions

View file

@ -10,7 +10,7 @@ use crate::{
};
use bevy_ptr::{OwningPtr, Ptr};
use bevy_utils::tracing::debug;
use std::any::TypeId;
use std::{any::TypeId, marker::PhantomData};
use super::{unsafe_world_cell::UnsafeEntityCell, Ref};
@ -1032,6 +1032,391 @@ impl<'w> EntityWorldMut<'w> {
pub fn update_location(&mut self) {
self.location = self.world.entities().get(self.entity).unwrap();
}
/// Gets an Entry into the world for this entity and component for in-place manipulation.
///
/// The type parameter specifies which component to get.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn_empty();
/// entity.entry().or_insert_with(|| Comp(4));
/// # let entity_id = entity.id();
/// assert_eq!(world.query::<&Comp>().single(&world).0, 4);
///
/// # let mut entity = world.get_entity_mut(entity_id).unwrap();
/// entity.entry::<Comp>().and_modify(|mut c| c.0 += 1);
/// assert_eq!(world.query::<&Comp>().single(&world).0, 5);
///
/// ```
pub fn entry<'a, T: Component>(&'a mut self) -> Entry<'w, 'a, T> {
if self.contains::<T>() {
Entry::Occupied(OccupiedEntry {
entity_world: self,
_marker: PhantomData,
})
} else {
Entry::Vacant(VacantEntry {
entity_world: self,
_marker: PhantomData,
})
}
}
}
/// A view into a single entity and component in a world, which may either be vacant or occupied.
///
/// This `enum` can only be constructed from the [`entry`] method on [`EntityWorldMut`].
///
/// [`entry`]: EntityWorldMut::entry
pub enum Entry<'w, 'a, T: Component> {
/// An occupied entry.
Occupied(OccupiedEntry<'w, 'a, T>),
/// A vacant entry.
Vacant(VacantEntry<'w, 'a, T>),
}
impl<'w, 'a, T: Component> Entry<'w, 'a, T> {
/// Provides in-place mutable access to an occupied entry.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn(Comp(0));
///
/// entity.entry::<Comp>().and_modify(|mut c| c.0 += 1);
/// assert_eq!(world.query::<&Comp>().single(&world).0, 1);
/// ```
#[inline]
pub fn and_modify<F: FnOnce(Mut<'_, T>)>(self, f: F) -> Self {
match self {
Entry::Occupied(mut entry) => {
f(entry.get_mut());
Entry::Occupied(entry)
}
Entry::Vacant(entry) => Entry::Vacant(entry),
}
}
/// Replaces the component of the entry, and returns an [`OccupiedEntry`].
///
/// # Examples
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn_empty();
///
/// let entry = entity.entry().insert_entry(Comp(4));
/// assert_eq!(entry.get(), &Comp(4));
///
/// let entry = entity.entry().insert_entry(Comp(2));
/// assert_eq!(entry.get(), &Comp(2));
/// ```
#[inline]
pub fn insert_entry(self, component: T) -> OccupiedEntry<'w, 'a, T> {
match self {
Entry::Occupied(mut entry) => {
entry.insert(component);
entry
}
Entry::Vacant(entry) => entry.insert_entry(component),
}
}
/// Ensures the entry has this component by inserting the given default if empty, and
/// returns a mutable reference to this component in the entry.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn_empty();
///
/// entity.entry().or_insert(Comp(4));
/// # let entity_id = entity.id();
/// assert_eq!(world.query::<&Comp>().single(&world).0, 4);
///
/// # let mut entity = world.get_entity_mut(entity_id).unwrap();
/// entity.entry().or_insert(Comp(15)).0 *= 2;
/// assert_eq!(world.query::<&Comp>().single(&world).0, 8);
/// ```
#[inline]
pub fn or_insert(self, default: T) -> Mut<'a, T> {
match self {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(default),
}
}
/// Ensures the entry has this component by inserting the result of the default function if
/// empty, and returns a mutable reference to this component in the entry.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn_empty();
///
/// entity.entry().or_insert_with(|| Comp(4));
/// assert_eq!(world.query::<&Comp>().single(&world).0, 4);
/// ```
#[inline]
pub fn or_insert_with<F: FnOnce() -> T>(self, default: F) -> Mut<'a, T> {
match self {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(default()),
}
}
}
impl<'w, 'a, T: Component + Default> Entry<'w, 'a, T> {
/// Ensures the entry has this component by inserting the default value if empty, and
/// returns a mutable reference to this component in the entry.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::prelude::*;
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn_empty();
///
/// entity.entry::<Comp>().or_default();
/// assert_eq!(world.query::<&Comp>().single(&world).0, 0);
/// ```
#[inline]
pub fn or_default(self) -> Mut<'a, T> {
match self {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert(Default::default()),
}
}
}
/// A view into an occupied entry in a [`EntityWorldMut`]. It is part of the [`Entry`] enum.
///
/// The contained entity must have the component type parameter if we have this struct.
pub struct OccupiedEntry<'w, 'a, T: Component> {
entity_world: &'a mut EntityWorldMut<'w>,
_marker: PhantomData<T>,
}
impl<'w, 'a, T: Component> OccupiedEntry<'w, 'a, T> {
/// Gets a reference to the component in the entry.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::{prelude::*, world::Entry};
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn(Comp(5));
///
/// if let Entry::Occupied(o) = entity.entry::<Comp>() {
/// assert_eq!(o.get().0, 5);
/// }
/// ```
#[inline]
pub fn get(&self) -> &T {
// This shouldn't panic because if we have an OccupiedEntry the component must exist.
self.entity_world.get::<T>().unwrap()
}
/// Gets a mutable reference to the component in the entry.
///
/// If you need a reference to the `OccupiedEntry` which may outlive the destruction of
/// the `Entry` value, see [`into_mut`].
///
/// [`into_mut`]: Self::into_mut
///
/// # Examples
///
/// ```
/// # use bevy_ecs::{prelude::*, world::Entry};
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn(Comp(5));
///
/// if let Entry::Occupied(mut o) = entity.entry::<Comp>() {
/// o.get_mut().0 += 10;
/// assert_eq!(o.get().0, 15);
///
/// // We can use the same Entry multiple times.
/// o.get_mut().0 += 2
/// }
///
/// assert_eq!(world.query::<&Comp>().single(&world).0, 17);
/// ```
#[inline]
pub fn get_mut(&mut self) -> Mut<'_, T> {
// This shouldn't panic because if we have an OccupiedEntry the component must exist.
self.entity_world.get_mut::<T>().unwrap()
}
/// Converts the `OccupiedEntry` into a mutable reference to the value in the entry with
/// a lifetime bound to the `EntityWorldMut`.
///
/// If you need multiple references to the `OccupiedEntry`, see [`get_mut`].
///
/// [`get_mut`]: Self::get_mut
///
/// # Examples
///
/// ```
/// # use bevy_ecs::{prelude::*, world::Entry};
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn(Comp(5));
///
/// if let Entry::Occupied(o) = entity.entry::<Comp>() {
/// o.into_mut().0 += 10;
/// }
///
/// assert_eq!(world.query::<&Comp>().single(&world).0, 15);
/// ```
#[inline]
pub fn into_mut(self) -> Mut<'a, T> {
// This shouldn't panic because if we have an OccupiedEntry the component must exist.
self.entity_world.get_mut().unwrap()
}
/// Replaces the component of the entry.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::{prelude::*, world::Entry};
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn(Comp(5));
///
/// if let Entry::Occupied(mut o) = entity.entry::<Comp>() {
/// o.insert(Comp(10));
/// }
///
/// assert_eq!(world.query::<&Comp>().single(&world).0, 10);
/// ```
#[inline]
pub fn insert(&mut self, component: T) {
self.entity_world.insert(component);
}
/// Removes the component from the entry and returns it.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::{prelude::*, world::Entry};
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn(Comp(5));
///
/// if let Entry::Occupied(o) = entity.entry::<Comp>() {
/// assert_eq!(o.take(), Comp(5));
/// }
///
/// assert_eq!(world.query::<&Comp>().iter(&world).len(), 0);
/// ```
#[inline]
pub fn take(self) -> T {
// This shouldn't panic because if we have an OccupiedEntry the component must exist.
self.entity_world.take().unwrap()
}
}
/// A view into a vacant entry in a [`EntityWorldMut`]. It is part of the [`Entry`] enum.
pub struct VacantEntry<'w, 'a, T: Component> {
entity_world: &'a mut EntityWorldMut<'w>,
_marker: PhantomData<T>,
}
impl<'w, 'a, T: Component> VacantEntry<'w, 'a, T> {
/// Inserts the component into the `VacantEntry` and returns a mutable reference to it.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::{prelude::*, world::Entry};
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn_empty();
///
/// if let Entry::Vacant(v) = entity.entry::<Comp>() {
/// v.insert(Comp(10));
/// }
///
/// assert_eq!(world.query::<&Comp>().single(&world).0, 10);
/// ```
#[inline]
pub fn insert(self, component: T) -> Mut<'a, T> {
self.entity_world.insert(component);
// This shouldn't panic because we just added this component
self.entity_world.get_mut::<T>().unwrap()
}
/// Inserts the component into the `VacantEntry` and returns an `OccupiedEntry`.
///
/// # Examples
///
/// ```
/// # use bevy_ecs::{prelude::*, world::Entry};
/// #[derive(Component, Default, Clone, Copy, Debug, PartialEq)]
/// struct Comp(u32);
///
/// # let mut world = World::new();
/// let mut entity = world.spawn_empty();
///
/// if let Entry::Vacant(v) = entity.entry::<Comp>() {
/// v.insert_entry(Comp(10));
/// }
///
/// assert_eq!(world.query::<&Comp>().single(&world).0, 10);
/// ```
#[inline]
pub fn insert_entry(self, component: T) -> OccupiedEntry<'w, 'a, T> {
self.entity_world.insert(component);
OccupiedEntry {
entity_world: self.entity_world,
_marker: PhantomData,
}
}
}
/// Inserts a dynamic [`Bundle`] into the entity.

View file

@ -7,7 +7,7 @@ pub mod unsafe_world_cell;
mod world_cell;
pub use crate::change_detection::{Mut, Ref, CHECK_TICK_THRESHOLD};
pub use entity_ref::{EntityMut, EntityRef, EntityWorldMut};
pub use entity_ref::{EntityMut, EntityRef, EntityWorldMut, Entry, OccupiedEntry, VacantEntry};
pub use spawn_batch::*;
pub use world_cell::*;