feat: add insert_if_new (#14397)

crates/bevy_ecs/src/archetype.rs

@@ -112,7 +112,7 @@ impl ArchetypeId {
 #[derive(Copy, Clone, Eq, PartialEq)]
 pub(crate) enum ComponentStatus {
     Added,
-    Mutated,
+    Existing,
 }
 
 pub(crate) struct AddBundle {
@@ -122,7 +122,7 @@ pub(crate) struct AddBundle {
     /// indicate if the component is newly added to the target archetype or if it already existed
     pub bundle_status: Vec<ComponentStatus>,
     pub added: Vec<ComponentId>,
-    pub mutated: Vec<ComponentId>,
+    pub existing: Vec<ComponentId>,
 }
 
 /// This trait is used to report the status of [`Bundle`](crate::bundle::Bundle) components
@@ -207,15 +207,15 @@ impl Edges {
         archetype_id: ArchetypeId,
         bundle_status: Vec<ComponentStatus>,
         added: Vec<ComponentId>,
-        mutated: Vec<ComponentId>,
+        existing: Vec<ComponentId>,
     ) {
         self.add_bundle.insert(
             bundle_id,
             AddBundle {
                 archetype_id,
                 bundle_status,
                 added,
-                mutated,
+                existing,
             },
         );
     }

crates/bevy_ecs/src/bundle.rs

@@ -300,6 +300,13 @@ impl SparseSetIndex for BundleId {
     }
 }
 
+// What to do on insertion if component already exists
+#[derive(Clone, Copy, Eq, PartialEq)]
+pub(crate) enum InsertMode {
+    Replace,
+    Keep,
+}
+
 /// Stores metadata associated with a specific type of [`Bundle`] for a given [`World`].
 ///
 /// [`World`]: crate::world::World
@@ -403,6 +410,7 @@ impl BundleInfo {
         table_row: TableRow,
         change_tick: Tick,
         bundle: T,
+        insert_mode: InsertMode,
         #[cfg(feature = "track_change_detection")] caller: &'static Location<'static>,
     ) {
         // NOTE: get_components calls this closure on each component in "bundle order".
@@ -418,8 +426,8 @@ impl BundleInfo {
                         unsafe { table.get_column_mut(component_id).debug_checked_unwrap() };
                     // SAFETY: bundle_component is a valid index for this bundle
                     let status = unsafe { bundle_component_status.get_status(bundle_component) };
-                    match status {
-                        ComponentStatus::Added => {
+                    match (status, insert_mode) {
+                        (ComponentStatus::Added, _) => {
                             column.initialize(
                                 table_row,
                                 component_ptr,
@@ -428,7 +436,7 @@ impl BundleInfo {
                                 caller,
                             );
                         }
-                        ComponentStatus::Mutated => {
+                        (ComponentStatus::Existing, InsertMode::Replace) => {
                             column.replace(
                                 table_row,
                                 component_ptr,
@@ -437,6 +445,9 @@ impl BundleInfo {
                                 caller,
                             );
                         }
+                        (ComponentStatus::Existing, InsertMode::Keep) => {
+                            column.drop(component_ptr);
+                        }
                     }
                 }
                 StorageType::SparseSet => {
@@ -482,7 +493,7 @@ impl BundleInfo {
         let current_archetype = &mut archetypes[archetype_id];
         for component_id in self.component_ids.iter().cloned() {
             if current_archetype.contains(component_id) {
-                bundle_status.push(ComponentStatus::Mutated);
+                bundle_status.push(ComponentStatus::Existing);
                 mutated.push(component_id);
             } else {
                 bundle_status.push(ComponentStatus::Added);
@@ -685,6 +696,7 @@ impl<'w> BundleInserter<'w> {
         entity: Entity,
         location: EntityLocation,
         bundle: T,
+        insert_mode: InsertMode,
         #[cfg(feature = "track_change_detection")] caller: &'static core::panic::Location<'static>,
     ) -> EntityLocation {
         let bundle_info = self.bundle_info.as_ref();
@@ -698,13 +710,15 @@ impl<'w> BundleInserter<'w> {
             // SAFETY: Mutable references do not alias and will be dropped after this block
             let mut deferred_world = self.world.into_deferred();
 
-            deferred_world.trigger_on_replace(
-                archetype,
-                entity,
-                add_bundle.mutated.iter().copied(),
-            );
-            if archetype.has_replace_observer() {
-                deferred_world.trigger_observers(ON_REPLACE, entity, &add_bundle.mutated);
+            if insert_mode == InsertMode::Replace {
+                deferred_world.trigger_on_replace(
+                    archetype,
+                    entity,
+                    add_bundle.existing.iter().copied(),
+                );
+                if archetype.has_replace_observer() {
+                    deferred_world.trigger_observers(ON_REPLACE, entity, &add_bundle.existing);
+                }
             }
         }
 
@@ -728,6 +742,7 @@ impl<'w> BundleInserter<'w> {
                     location.table_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
                     #[cfg(feature = "track_change_detection")]
                     caller,
                 );
@@ -768,6 +783,7 @@ impl<'w> BundleInserter<'w> {
                     result.table_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
                     #[cfg(feature = "track_change_detection")]
                     caller,
                 );
@@ -849,6 +865,7 @@ impl<'w> BundleInserter<'w> {
                     move_result.new_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
                     #[cfg(feature = "track_change_detection")]
                     caller,
                 );
@@ -970,6 +987,7 @@ impl<'w> BundleSpawner<'w> {
                 table_row,
                 self.change_tick,
                 bundle,
+                InsertMode::Replace,
                 #[cfg(feature = "track_change_detection")]
                 caller,
             );
@@ -1223,6 +1241,9 @@ mod tests {
     #[derive(Component)]
     struct D;
 
+    #[derive(Component, Eq, PartialEq, Debug)]
+    struct V(&'static str); // component with a value
+
     #[derive(Resource, Default)]
     struct R(usize);
 
@@ -1295,6 +1316,7 @@ mod tests {
         world.init_resource::<R>();
         let mut entity = world.entity_mut(entity);
         entity.insert(A);
+        entity.insert_if_new(A); // this will not trigger on_replace or on_insert
         entity.flush();
         assert_eq!(2, world.resource::<R>().0);
     }
@@ -1364,4 +1386,19 @@ mod tests {
         world.spawn(A).flush();
         assert_eq!(4, world.resource::<R>().0);
     }
+
+    #[test]
+    fn insert_if_new() {
+        let mut world = World::new();
+        world.init_resource::<R>();
+        let id = world.spawn(V("one")).id();
+        let mut entity = world.entity_mut(id);
+        entity.insert_if_new(V("two"));
+        entity.insert_if_new((A, V("three")));
+        entity.flush();
+        // should still contain "one"
+        let entity = world.entity(id);
+        assert!(entity.contains::<A>());
+        assert_eq!(entity.get(), Some(&V("one")));
+    }
 }

crates/bevy_ecs/src/storage/blob_vec.rs

@@ -437,6 +437,12 @@ impl BlobVec {
             }
         }
     }
+
+    /// Get the dropper for this vector. Used if caller has type-erased pointer
+    /// which they have decided _not_ to add to this vector.
+    pub fn get_drop(&self) -> Option<unsafe fn(OwningPtr<'_>)> {
+        self.drop
+    }
 }
 
 impl Drop for BlobVec {

crates/bevy_ecs/src/storage/table.rs

@@ -232,6 +232,20 @@ impl Column {
         }
     }
 
+    /// Call [`drop`] on a value.
+    ///
+    /// # Safety
+    /// [`data`] must point to the same type that this table stores, so the
+    /// correct drop function is called.
+    #[inline]
+    pub(crate) unsafe fn drop(&self, data: OwningPtr<'_>) {
+        if let Some(drop) = self.data.get_drop() {
+            // Safety: we're using the same drop fn that the BlobVec would
+            // if we inserted the data instead of dropping it.
+            unsafe { drop(data) }
+        }
+    }
+
     /// Gets the current number of elements stored in the column.
     #[inline]
     pub fn len(&self) -> usize {

crates/bevy_ecs/src/system/commands/mod.rs

@@ -5,7 +5,7 @@ use core::panic::Location;
 use super::{Deferred, IntoObserverSystem, IntoSystem, RegisterSystem, Resource};
 use crate::{
     self as bevy_ecs,
-    bundle::Bundle,
+    bundle::{Bundle, InsertMode},
     component::{ComponentId, ComponentInfo},
     entity::{Entities, Entity},
     event::Event,
@@ -895,6 +895,7 @@ impl EntityCommands<'_> {
     /// Adds a [`Bundle`] of components to the entity.
     ///
     /// This will overwrite any previous value(s) of the same component type.
+    /// See [`EntityCommands::insert_if_new`] to keep the old value instead.
     ///
     /// # Panics
     ///
@@ -945,7 +946,22 @@ impl EntityCommands<'_> {
     /// ```
     #[track_caller]
     pub fn insert(&mut self, bundle: impl Bundle) -> &mut Self {
-        self.add(insert(bundle))
+        self.add(insert(bundle, InsertMode::Replace))
+    }
+
+    /// Adds a [`Bundle`] of components to the entity.
+    ///
+    /// This is the same as [`insert`], but in case of duplicate components
+    /// will leave the old values instead of replacing them with new ones.
+    ///
+    /// # Panics
+    ///
+    /// The command will panic when applied if the associated entity does not exist.
+    ///
+    /// To avoid a panic in this case, use the command [`Self::try_insert`] instead.
+    /// ```
+    pub fn insert_if_new(&mut self, bundle: impl Bundle) -> &mut Self {
+        self.add(insert(bundle, InsertMode::Keep))
     }
 
     /// Adds a dynamic component to an entity.
@@ -1044,7 +1060,7 @@ impl EntityCommands<'_> {
     /// ```
     #[track_caller]
     pub fn try_insert(&mut self, bundle: impl Bundle) -> &mut Self {
-        self.add(try_insert(bundle))
+        self.add(try_insert(bundle, InsertMode::Replace))
     }
 
     /// Removes a [`Bundle`] of components from the entity.
@@ -1312,12 +1328,13 @@ fn despawn() -> impl EntityCommand {
 
 /// An [`EntityCommand`] that adds the components in a [`Bundle`] to an entity.
 #[track_caller]
-fn insert<T: Bundle>(bundle: T) -> impl EntityCommand {
+fn insert<T: Bundle>(bundle: T, mode: InsertMode) -> impl EntityCommand {
     let caller = core::panic::Location::caller();
     move |entity: Entity, world: &mut World| {
         if let Some(mut entity) = world.get_entity_mut(entity) {
             entity.insert_with_caller(
                 bundle,
+                mode,
                 #[cfg(feature = "track_change_detection")]
                 caller,
             );
@@ -1328,14 +1345,16 @@ fn insert<T: Bundle>(bundle: T) -> impl EntityCommand {
 }
 
 /// An [`EntityCommand`] that attempts to add the components in a [`Bundle`] to an entity.
+/// Does nothing if the entity does not exist.
 #[track_caller]
-fn try_insert(bundle: impl Bundle) -> impl EntityCommand {
+fn try_insert(bundle: impl Bundle, mode: InsertMode) -> impl EntityCommand {
     #[cfg(feature = "track_change_detection")]
     let caller = core::panic::Location::caller();
     move |entity, world: &mut World| {
         if let Some(mut entity) = world.get_entity_mut(entity) {
             entity.insert_with_caller(
                 bundle,
+                mode,
                 #[cfg(feature = "track_change_detection")]
                 caller,
             );

crates/bevy_ecs/src/world/entity_ref.rs

@@ -1,6 +1,6 @@
 use crate::{
     archetype::{Archetype, ArchetypeId, Archetypes},
-    bundle::{Bundle, BundleId, BundleInfo, BundleInserter, DynamicBundle},
+    bundle::{Bundle, BundleId, BundleInfo, BundleInserter, DynamicBundle, InsertMode},
     change_detection::MutUntyped,
     component::{Component, ComponentId, ComponentTicks, Components, StorageType},
     entity::{Entities, Entity, EntityLocation},
@@ -772,6 +772,20 @@ impl<'w> EntityWorldMut<'w> {
     pub fn insert<T: Bundle>(&mut self, bundle: T) -> &mut Self {
         self.insert_with_caller(
             bundle,
+            InsertMode::Replace,
+            #[cfg(feature = "track_change_detection")]
+            core::panic::Location::caller(),
+        )
+    }
+
+    /// Adds a [`Bundle`] of components to the entity.
+    ///
+    /// This will overwrite any previous value(s) of the same component type.
+    #[track_caller]
+    pub fn insert_if_new<T: Bundle>(&mut self, bundle: T) -> &mut Self {
+        self.insert_with_caller(
+            bundle,
+            InsertMode::Keep,
             #[cfg(feature = "track_change_detection")]
             core::panic::Location::caller(),
         )
@@ -783,6 +797,7 @@ impl<'w> EntityWorldMut<'w> {
     pub(crate) fn insert_with_caller<T: Bundle>(
         &mut self,
         bundle: T,
+        mode: InsertMode,
         #[cfg(feature = "track_change_detection")] caller: &'static core::panic::Location,
     ) -> &mut Self {
         let change_tick = self.world.change_tick();
@@ -791,7 +806,7 @@ impl<'w> EntityWorldMut<'w> {
         self.location =
             // SAFETY: location matches current entity. `T` matches `bundle_info`
             unsafe {
-                bundle_inserter.insert(self.entity, self.location, bundle,  #[cfg(feature = "track_change_detection")] caller)
+                bundle_inserter.insert(self.entity, self.location, bundle, mode, #[cfg(feature = "track_change_detection")] caller)
             };
         self
     }
@@ -2355,6 +2370,7 @@ unsafe fn insert_dynamic_bundle<
             entity,
             location,
             bundle,
+            InsertMode::Replace,
             #[cfg(feature = "track_change_detection")]
             core::panic::Location::caller(),
         )