feat: add insert_if_new (#14397) (#14646)

# Objective

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.

See also issue #14397

Fixes #2054.

## Solution

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.

## Testing

*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.

*How can other people (reviewers) test your changes? Is there anything
specific they need to know?*

`cargo test` in the bevy_ecs project.

*If relevant, what platforms did you test these changes on, and are
there any important ones you can't test?*

Only tested on Windows, but it doesn't touch anything platform-specific.

---------

Co-authored-by: Alice Cecile <[email protected]>
Co-authored-by: Giacomo Stevanato <[email protected]>

Author: 3 people

crates/bevy_ecs/src/archetype.rs

@@ -109,10 +109,12 @@ impl ArchetypeId {
     }
 }
 
+/// Used in [`AddBundle`] to track whether components in the bundle are newly
+/// added or already existed in the entity's archetype.
 #[derive(Copy, Clone, Eq, PartialEq)]
 pub(crate) enum ComponentStatus {
     Added,
-    Mutated,
+    Existing,
 }
 
 pub(crate) struct AddBundle {
@@ -122,7 +124,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 +209,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

@@ -307,6 +307,15 @@ impl SparseSetIndex for BundleId {
     }
 }
 
+// What to do on insertion if component already exists
+#[derive(Clone, Copy, Eq, PartialEq)]
+pub(crate) enum InsertMode {
+    /// Any existing components of a matching type will be overwritten.
+    Replace,
+    /// Any existing components of a matching type will kept unchanged.
+    Keep,
+}
+
 /// Stores metadata associated with a specific type of [`Bundle`] for a given [`World`].
 ///
 /// [`World`]: crate::world::World
@@ -410,6 +419,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".
@@ -425,8 +435,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,
@@ -435,7 +445,7 @@ impl BundleInfo {
                                 caller,
                             );
                         }
-                        ComponentStatus::Mutated => {
+                        (ComponentStatus::Existing, InsertMode::Replace) => {
                             column.replace(
                                 table_row,
                                 component_ptr,
@@ -444,6 +454,9 @@ impl BundleInfo {
                                 caller,
                             );
                         }
+                        (ComponentStatus::Existing, InsertMode::Keep) => {
+                            column.drop(component_ptr);
+                        }
                     }
                 }
                 StorageType::SparseSet => {
@@ -489,7 +502,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);
@@ -692,6 +705,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();
@@ -705,13 +719,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);
+                }
             }
         }
 
@@ -735,6 +751,7 @@ impl<'w> BundleInserter<'w> {
                     location.table_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
                     #[cfg(feature = "track_change_detection")]
                     caller,
                 );
@@ -775,6 +792,7 @@ impl<'w> BundleInserter<'w> {
                     result.table_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
                     #[cfg(feature = "track_change_detection")]
                     caller,
                 );
@@ -856,6 +874,7 @@ impl<'w> BundleInserter<'w> {
                     move_result.new_row,
                     self.change_tick,
                     bundle,
+                    insert_mode,
                     #[cfg(feature = "track_change_detection")]
                     caller,
                 );
@@ -875,9 +894,34 @@ impl<'w> BundleInserter<'w> {
             if new_archetype.has_add_observer() {
                 deferred_world.trigger_observers(ON_ADD, entity, &add_bundle.added);
             }
-            deferred_world.trigger_on_insert(new_archetype, entity, bundle_info.iter_components());
-            if new_archetype.has_insert_observer() {
-                deferred_world.trigger_observers(ON_INSERT, entity, bundle_info.components());
+            match insert_mode {
+                InsertMode::Replace => {
+                    // insert triggers for both new and existing components if we're replacing them
+                    deferred_world.trigger_on_insert(
+                        new_archetype,
+                        entity,
+                        bundle_info.iter_components(),
+                    );
+                    if new_archetype.has_insert_observer() {
+                        deferred_world.trigger_observers(
+                            ON_INSERT,
+                            entity,
+                            bundle_info.components(),
+                        );
+                    }
+                }
+                InsertMode::Keep => {
+                    // insert triggers only for new components if we're not replacing them (since
+                    // nothing is actually inserted).
+                    deferred_world.trigger_on_insert(
+                        new_archetype,
+                        entity,
+                        add_bundle.added.iter().cloned(),
+                    );
+                    if new_archetype.has_insert_observer() {
+                        deferred_world.trigger_observers(ON_INSERT, entity, &add_bundle.added);
+                    }
+                }
             }
         }
 
@@ -977,6 +1021,7 @@ impl<'w> BundleSpawner<'w> {
                 table_row,
                 self.change_tick,
                 bundle,
+                InsertMode::Replace,
                 #[cfg(feature = "track_change_detection")]
                 caller,
             );
@@ -1230,6 +1275,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);
 
@@ -1302,6 +1350,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);
     }
@@ -1371,4 +1420,18 @@ mod tests {
         world.spawn(A).flush();
         assert_eq!(4, world.resource::<R>().0);
     }
+
+    #[test]
+    fn insert_if_new() {
+        let mut world = World::new();
+        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,14 @@ impl BlobVec {
             }
         }
     }
+
+    /// Get the `drop` argument that was passed to `BlobVec::new`.
+    ///
+    /// Callers can use this if they have a type-erased pointer of the correct
+    /// type to add to this [`BlobVec`], which they just want to drop instead.
+    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,24 @@ 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 without overwriting.
+    ///
+    /// This is the same as [`EntityCommands::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 +1062,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.
@@ -1321,12 +1339,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,
             );
@@ -1337,14 +1356,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,
             );