diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 274daeed05547..c94b45dad10ee 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 0919dfd6324c6..0de995257fc4d 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::*;