From 865041de74da409603c7d413313a607bed68c3ec Mon Sep 17 00:00:00 2001 From: TheBigCheese <32036861+13ros27@users.noreply.github.com> Date: Tue, 21 Nov 2023 01:26:06 +0000 Subject: [PATCH] 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 Co-authored-by: Pascal Hertleif --- crates/bevy_ecs/src/world/entity_ref.rs | 387 +++++++++++++++++++++++- crates/bevy_ecs/src/world/mod.rs | 2 +- 2 files changed, 387 insertions(+), 2 deletions(-) diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 274daeed05..c94b45dad1 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -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::().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::() { + 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::().and_modify(|mut c| c.0 += 1); + /// assert_eq!(world.query::<&Comp>().single(&world).0, 1); + /// ``` + #[inline] + pub fn and_modify)>(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 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::().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, +} + +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::() { + /// 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::().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::() { + /// 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::().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::() { + /// 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::() { + /// 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::() { + /// 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, +} + +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::() { + /// 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::().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::() { + /// 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. diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index 0919dfd632..0de995257f 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -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::*;