Add an Entry api to EntityWorldMut.

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,388 @@ 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(Comp(0));
+    /// entity.entry::<Comp>().and_modify(|mut c| c.0 += 1);
+    /// # assert_eq!(world.query::<&Comp>().single(&world).0, 1);
+    ///
+    /// let mut new_entity = world.spawn_empty();
+    /// entity.entry().or_insert_with(|| Comp(4));
+    /// # assert_eq!(world.query::<&Comp>().single(&world).0, 4);
+    /// ```
+    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` is 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.
+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 {
+        // SAFETY: If we have an OccupiedEntry the component must exist.
+        unsafe { self.entity_world.get::<T>().unwrap_unchecked() }
+    }
+
+    /// 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> {
+        // SAFETY: If we have an OccupiedEntry the component must exist.
+        unsafe { self.entity_world.get_mut::<T>().unwrap_unchecked() }
+    }
+
+    /// 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> {
+        // SAFETY: If we have an OccupiedEntry the component must exist.
+        unsafe { self.entity_world.get_mut().unwrap_unchecked() }
+    }
+
+    /// 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 {
+        // SAFETY: If we have an OccupiedEntry the component must exist.
+        unsafe { self.entity_world.take().unwrap_unchecked() }
+    }
+}
+
+/// 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);
+        // SAFETY: We just added this component.
+        unsafe { self.entity_world.get_mut::<T>().unwrap_unchecked() }
+    }
+
+    /// 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.