diff --git a/crates/bevy_ecs/src/system/commands.rs b/crates/bevy_ecs/src/system/commands.rs index 90adfb8dc510c..e2308fa9db6fa 100644 --- a/crates/bevy_ecs/src/system/commands.rs +++ b/crates/bevy_ecs/src/system/commands.rs @@ -207,6 +207,7 @@ impl<'a, 'b> EntityCommands<'a, 'b> { } /// Adds a single [`Component`] to the current entity. + /// This will overwrite an existing [`Component`] of the same type. /// /// /// # Warning @@ -246,6 +247,16 @@ impl<'a, 'b> EntityCommands<'a, 'b> { self } + /// Adds a single [`Component`] to the current entity, if (and only if) the current entity + /// does not already have a [`Component`] of the same type. + pub fn try_insert(&mut self, component: impl Component) -> &mut Self { + self.commands.add(TryInsert { + entity: self.entity, + component, + }); + self + } + /// See [`EntityMut::remove_bundle`](crate::world::EntityMut::remove_bundle). pub fn remove_bundle(&mut self) -> &mut Self where @@ -357,6 +368,24 @@ where } } +#[derive(Debug)] +pub struct TryInsert { + pub entity: Entity, + pub component: T, +} + +impl Command for TryInsert +where + T: Component, +{ + fn write(self: Box, world: &mut World) { + world + .entity_mut(self.entity) + .component::() + .or_insert(self.component); + } +} + #[derive(Debug)] pub struct Remove { entity: Entity, diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 1815d15798f43..7185045c3d581 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -7,7 +7,7 @@ use crate::{ storage::{SparseSet, Storages}, world::{Mut, World}, }; -use std::any::TypeId; +use std::{any::TypeId, marker::PhantomData}; pub struct EntityRef<'w> { world: &'w World, @@ -461,6 +461,21 @@ impl<'w> EntityMut<'w> { entities.meta[self.entity.id as usize].location = new_location; } + /// Gets the current entity's [`ComponentEntry`] for in-place manipulation. + pub fn component(&mut self) -> ComponentEntry<'_, 'w, T> { + if let Some(component) = self.get_mut::() { + ComponentEntry::Occupied(OccupiedComponentEntry { + entity: self, + component, + }) + } else { + ComponentEntry::Vacant(VacantComponentEntry { + entity: self, + _phantom: PhantomData, + }) + } + } + pub fn insert(&mut self, value: T) -> &mut Self { self.insert_bundle((value,)) } @@ -877,6 +892,136 @@ fn sorted_remove(source: &mut Vec, remove: &[T]) { }) } +/// A view into an occupied entity's component entry. +/// It is part of the [`ComponentEntry`] enum. +pub struct OccupiedComponentEntry<'e, 'w, T> { + entity: &'e mut EntityMut<'w>, + component: Mut<'w, T>, +} + +impl<'e, 'w, T> OccupiedComponentEntry<'e, 'w, T> +where + T: Component, +{ + /// Gets a reference to the entity's component. + #[inline] + pub fn get(&self) -> &T { + self.component.as_ref() + } + + /// Gets a mutable reference to the entity's component. + /// + /// ## Note + /// This triggers change detection. + #[inline] + pub fn get_mut(&mut self) -> &mut T { + self.component.as_mut() + } + + /// Sets the value of the component, and returns the old component. + #[inline] + pub fn insert(&mut self, component: T) -> T { + std::mem::replace(self.component.as_mut(), component) + } + + /// Gets a [`Mut`] for the entity's component. + #[inline] + pub fn into_mut(self) -> Mut<'w, T> { + self.component + } + + /// Removes the entity's component, and returns it. + #[inline] + pub fn remove(self) -> T { + self.entity.remove::().unwrap() + } +} + +/// A view into a vacant entity's component entry. +/// It is part of the [`ComponentEntry`] enum. +pub struct VacantComponentEntry<'e, 'w, T> { + entity: &'e mut EntityMut<'w>, + _phantom: PhantomData, +} + +impl<'e, 'w, T> VacantComponentEntry<'e, 'w, T> +where + T: Component, +{ + /// Inserts the component for the current entity, and returns + /// a [`Mut`] of the component. + #[inline] + pub fn insert(self, component: T) -> Mut<'w, T> { + self.entity.insert(component); + self.entity.get_mut().unwrap() + } +} + +/// A view into a single entity's component in the [`World`], which may either +/// be vacant or occupied. +/// This `enum` is constructed from the [`component`] method on [`EntityMut`] +/// +/// This is analogous to [`entry`] for [`HashMap`] +/// +/// [`entry`]: std::collections::hash_map::Entry +/// [`HashMap`]: std::collections::HashMap +/// [`component`]: EntityMut::component +pub enum ComponentEntry<'e, 'w, T: Component> { + /// An occupied component entry.. + Occupied(OccupiedComponentEntry<'e, 'w, T>), + /// A vacant component entry. + Vacant(VacantComponentEntry<'e, 'w, T>), +} + +impl<'e, 'w, T: Component> ComponentEntry<'e, 'w, T> { + /// Provides in-place mutable access to the current entity's component + /// before any potential inserts into the world. + /// + /// ## Note + /// If the component exists, this triggers change detection. + #[inline] + pub fn and_modify(self, f: impl FnOnce(&mut T)) -> Self { + match self { + ComponentEntry::Occupied(mut o) => { + f(o.get_mut()); + ComponentEntry::Occupied(o) + } + ComponentEntry::Vacant(v) => ComponentEntry::Vacant(v), + } + } + + /// Ensures a component is in the entry by inserting the default if empty, and returns + /// a [`Mut`] to the component in the entry. + /// + /// The `component` passed to `or_insert` is eagerly evaluated; if you are passing the result of + /// a function call, it is recommended to use [`or_insert_with`], which is lazily evaluated. + /// + /// [`or_insert_with`]: ComponentEntry::or_insert_with + #[inline] + pub fn or_insert(self, component: T) -> Mut<'w, T> { + self.or_insert_with(|| component) + } + + /// Ensures a component is in the entry by inserting the result of the `default` function + /// if empty, and returns a [`Mut`] to the component in the entry. + #[inline] + pub fn or_insert_with(self, default: impl FnOnce() -> T) -> Mut<'w, T> { + match self { + ComponentEntry::Occupied(o) => o.component, + ComponentEntry::Vacant(v) => v.insert(default()), + } + } +} + +impl<'e, 'w, T: Component + Default> ComponentEntry<'e, 'w, T> { + /// Ensures a component is in the entry by inserting the default value if empty, and returns + /// a [`Mut`] to the component in the entry. + #[inline] + pub fn or_default(self) -> Mut<'w, T> { + self.or_insert_with(T::default) + } +} + #[cfg(test)] mod tests { #[test]