feat: add insert_if_new (bevyengine#14397)

Often there are reasons to insert some components (e.g. Transform)
separately from the rest of a bundle (e.g. PbrBundle).
However `insert` overwrites existing components, making this difficult.

This PR adds the method `insert_if_new` to EntityMut and Commands, which
is the same as `insert` except that the old component is kept in case of
conflicts.

It also renames some internal enums (from `ComponentStatus::Mutated` to
`Existing`), to reflect the possible change in meaning.

- Did you test these changes? If so, how?

Added basic unit tests; used the new behavior in my project.

- Are there any parts that need more testing?

There should be a test that the change time isn't set if a component is
not overwritten; I wasn't sure how to write a test for that case.

crates/bevy_ecs/src/archetype.rs

@@ -111,7 +111,7 @@ impl ArchetypeId {
 #[derive(Copy, Clone, Eq, PartialEq)]
 pub(crate) enum ComponentStatus {
     Added,
-    Mutated,
+    Existing,
 }
 
 pub(crate) struct AddBundle {

crates/bevy_ecs/src/bundle.rs

@@ -298,6 +298,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
@@ -401,6 +408,8 @@ 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".
         // bundle_info.component_ids are also in "bundle order"
@@ -415,12 +424,27 @@ 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 => {
-                            column.initialize(table_row, component_ptr, change_tick);
+                    match (status, insert_mode) {
+                        (ComponentStatus::Added, _) => {
+                            column.initialize(
+                                table_row,
+                                component_ptr,
+                                change_tick,
+                                #[cfg(feature = "track_change_detection")]
+                                caller,
+                            );
                         }
-                        ComponentStatus::Mutated => {
-                            column.replace(table_row, component_ptr, change_tick);
+                        (ComponentStatus::Existing, InsertMode::Replace) => {
+                            column.replace(
+                                table_row,
+                                component_ptr,
+                                change_tick,
+                                #[cfg(feature = "track_change_detection")]
+                                caller,
+                            );
+                        }
+                        (ComponentStatus::Existing, InsertMode::Keep) => {
+                            column.drop(component_ptr);
                         }
                     }
                 }
@@ -460,7 +484,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);
             } else {
                 bundle_status.push(ComponentStatus::Added);
                 added.push(component_id);
@@ -661,10 +685,15 @@ 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();
         let add_bundle = self.add_bundle.as_ref();
         let table = self.table.as_mut();
+
+        // SAFETY: Archetype gets borrowed when running the on_replace observers above,
+        // so this reference can only be promoted from shared to &mut down here, after they have been ran
         let archetype = self.archetype.as_mut();
 
         let (new_archetype, new_location) = match &mut self.result {
@@ -683,6 +712,9 @@ impl<'w> BundleInserter<'w> {
                     location.table_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
+                    #[cfg(feature = "track_change_detection")]
+                    caller,
                 );
 
                 (archetype, location)
@@ -721,6 +753,9 @@ impl<'w> BundleInserter<'w> {
                     result.table_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
+                    #[cfg(feature = "track_change_detection")]
+                    caller,
                 );
 
                 (new_archetype, new_location)
@@ -800,6 +835,9 @@ impl<'w> BundleInserter<'w> {
                     move_result.new_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
+                    #[cfg(feature = "track_change_detection")]
+                    caller,
                 );
 
                 (new_archetype, new_location)
@@ -918,6 +956,9 @@ impl<'w> BundleSpawner<'w> {
                 table_row,
                 self.change_tick,
                 bundle,
+                InsertMode::Replace,
+                #[cfg(feature = "track_change_detection")]
+                caller,
             );
             entities.set(entity.index(), location);
             location
@@ -1138,6 +1179,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);
 
@@ -1186,7 +1230,7 @@ mod tests {
         entity.insert(A);
         entity.remove::<A>();
         entity.flush();
-        assert_eq!(3, world.resource::<R>().0);
+        assert_eq!(4, world.resource::<R>().0);
     }
 
     #[test]
@@ -1254,4 +1298,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

@@ -204,6 +204,18 @@ impl Column {
             .get_mut() = change_tick;
     }
 
+    /// 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() {
+            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

@@ -3,7 +3,7 @@ mod parallel_scope;
 use super::{Deferred, IntoObserverSystem, IntoSystem, RegisterSystem, Resource};
 use crate::{
     self as bevy_ecs,
-    bundle::Bundle,
+    bundle::{Bundle, InsertMode},
     component::ComponentId,
     entity::{Entities, Entity},
     event::Event,
@@ -894,6 +894,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
     ///
@@ -943,7 +944,22 @@ impl EntityCommands<'_> {
     /// # bevy_ecs::system::assert_is_system(add_combat_stats_system);
     /// ```
     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))
     }
 
     /// Tries to add a [`Bundle`] of components to the entity.
@@ -995,7 +1011,7 @@ impl EntityCommands<'_> {
     /// # bevy_ecs::system::assert_is_system(add_combat_stats_system);
     /// ```
     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.
@@ -1242,21 +1258,37 @@ fn despawn(entity: Entity, world: &mut World) {
 }
 
 /// An [`EntityCommand`] that adds the components in a [`Bundle`] to an entity.
-fn insert<T: Bundle>(bundle: T) -> impl EntityCommand {
+#[track_caller]
+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(bundle);
+            entity.insert_with_caller(
+                bundle,
+                mode,
+                #[cfg(feature = "track_change_detection")]
+                caller,
+            );
         } else {
             panic!("error[B0003]: Could not insert a bundle (of type `{}`) for entity {:?} because it doesn't exist in this World. See: https://bevyengine.org/learn/errors/#b0003", std::any::type_name::<T>(), entity);
         }
     }
 }
 
 /// An [`EntityCommand`] that attempts to add the components in a [`Bundle`] to an entity.
-fn try_insert(bundle: impl Bundle) -> impl EntityCommand {
+/// Does nothing if the entity does not exist.
+#[track_caller]
+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(bundle);
+            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},
@@ -769,11 +769,44 @@ impl<'w> EntityWorldMut<'w> {
     ///
     /// This will overwrite any previous value(s) of the same component type.
     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(),
+        )
+    }
+
+    /// Split into a new function so we can pass the calling location into the function when using
+    /// as a command.
+    #[inline]
+    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();
         let mut bundle_inserter =
             BundleInserter::new::<T>(self.world, self.location.archetype_id, change_tick);
-        // SAFETY: location matches current entity. `T` matches `bundle_info`
-        self.location = unsafe { bundle_inserter.insert(self.entity, self.location, bundle) };
+        self.location =
+            // SAFETY: location matches current entity. `T` matches `bundle_info`
+            unsafe {
+                bundle_inserter.insert(self.entity, self.location, bundle, mode, #[cfg(feature = "track_change_detection")] caller)
+            };
         self
     }
 
@@ -2312,7 +2345,16 @@ unsafe fn insert_dynamic_bundle<
     };
 
     // SAFETY: location matches current entity.
-    unsafe { bundle_inserter.insert(entity, location, bundle) }
+    unsafe {
+        bundle_inserter.insert(
+            entity,
+            location,
+            bundle,
+            InsertMode::Replace,
+            #[cfg(feature = "track_change_detection")]
+            core::panic::Location::caller(),
+        )
+    }
 }
 
 /// Removes a bundle from the given archetype and returns the resulting archetype (or None if the