Add non-overwriting variants of component/resource insertion functions

crates/bevy_ecs/src/bundle.rs

@@ -166,6 +166,50 @@ impl BundleInfo {
         });
     }
 
+    /// Like write_components, but does not overwrite existing components on the entity
+    /// # Safety
+    /// table row must exist, entity must be valid
+    #[allow(clippy::clippy::too_many_arguments)]
+    #[inline]
+    pub(crate) unsafe fn write_additional_components<T: Bundle>(
+        &self,
+        sparse_sets: &mut SparseSets,
+        entity: Entity,
+        table: &Table,
+        table_row: usize,
+        bundle_status: &[ComponentStatus],
+        bundle: T,
+        change_tick: u32,
+    ) {
+        // NOTE: get_components calls this closure on each component in "bundle order".
+        // bundle_info.component_ids are also in "bundle order"
+        let mut bundle_component = 0;
+        bundle.get_components(|component_ptr| {
+            // SAFE: component_id was initialized by get_dynamic_bundle_info
+            let component_id = *self.component_ids.get_unchecked(bundle_component);
+            let component_status = bundle_status.get_unchecked(bundle_component);
+            match self.storage_types[bundle_component] {
+                StorageType::Table => match component_status {
+                    ComponentStatus::Added => {
+                        let column = table.get_column(component_id).unwrap();
+                        column.set_unchecked(table_row, component_ptr);
+                        let column_status = column.get_ticks_unchecked_mut(table_row);
+                        *column_status = ComponentTicks::new(change_tick);
+                    }
+                    ComponentStatus::Mutated => {}
+                },
+                StorageType::SparseSet => match component_status {
+                    ComponentStatus::Added => {
+                        let sparse_set = sparse_sets.get_mut(component_id).unwrap();
+                        sparse_set.insert(entity, component_ptr, change_tick);
+                    }
+                    ComponentStatus::Mutated => {}
+                },
+            }
+            bundle_component += 1;
+        });
+    }
+
     #[inline]
     pub fn id(&self) -> BundleId {
         self.id

crates/bevy_ecs/src/system/commands.rs

@@ -206,6 +206,15 @@ impl<'a, 'b> EntityCommands<'a, 'b> {
         self
     }
 
+    /// Adds a bundle of components to the current entity, skipping ones that already exist.
+    pub fn try_insert_bundle(&mut self, bundle: impl Bundle) -> &mut Self {
+        self.commands.add(TryInsertBundle {
+            entity: self.entity,
+            bundle,
+        });
+        self
+    }
+
     /// Adds a single [`Component`] to the current entity.
     ///
     ///
@@ -246,6 +255,15 @@ impl<'a, 'b> EntityCommands<'a, 'b> {
         self
     }
 
+    /// Like insert, but does not overwrite components which already exist on the entity
+    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<T>(&mut self) -> &mut Self
     where
@@ -342,6 +360,20 @@ where
     }
 }
 
+pub struct TryInsertBundle<T> {
+    entity: Entity,
+    bundle: T,
+}
+
+impl<T> Command for TryInsertBundle<T>
+where
+    T: Bundle + 'static,
+{
+    fn write(self: Box<Self>, world: &mut World) {
+        world.entity_mut(self.entity).try_insert_bundle(self.bundle);
+    }
+}
+
 #[derive(Debug)]
 pub struct Insert<T> {
     pub entity: Entity,
@@ -357,6 +389,21 @@ where
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct TryInsert<T> {
+    entity: Entity,
+    component: T,
+}
+
+impl<T> Command for TryInsert<T>
+where
+    T: Component,
+{
+    fn write(self: Box<Self>, world: &mut World) {
+        world.entity_mut(self.entity).try_insert(self.component);
+    }
+}
+
 #[derive(Debug)]
 pub struct Remove<T> {
     entity: Entity,
@@ -451,6 +498,64 @@ mod tests {
         assert_eq!(results2, vec![]);
     }
 
+    #[test]
+    fn try_insert_components_not_present() {
+        let mut world = World::default();
+        let mut command_queue = CommandQueue::default();
+        let _ = Commands::new(&mut command_queue, &world)
+            .spawn()
+            .try_insert_bundle((2u32, 4u64))
+            .id();
+
+        command_queue.apply(&mut world);
+        assert!(world.entities().len() == 1);
+        let results = world
+            .query::<(&u32, &u64)>()
+            .iter(&world)
+            .map(|(a, b)| (*a, *b))
+            .collect::<Vec<_>>();
+        assert_eq!(results, vec![(2u32, 4u64)]);
+    }
+
+    #[test]
+    fn try_insert_components_present() {
+        let mut world = World::default();
+        let mut command_queue = CommandQueue::default();
+        let _ = Commands::new(&mut command_queue, &world)
+            .spawn_bundle((1u32, 2u64))
+            .try_insert_bundle((2u32, 4u64))
+            .id();
+
+        command_queue.apply(&mut world);
+        assert!(world.entities().len() == 1);
+        let results = world
+            .query::<(&u32, &u64)>()
+            .iter(&world)
+            .map(|(a, b)| (*a, *b))
+            .collect::<Vec<_>>();
+        assert_eq!(results, vec![(1u32, 2u64)]);
+    }
+
+    #[test]
+    fn try_insert_components_some_present() {
+        let mut world = World::default();
+        let mut command_queue = CommandQueue::default();
+        let _ = Commands::new(&mut command_queue, &world)
+            .spawn()
+            .insert(1u32)
+            .try_insert_bundle((2u32, 4u64))
+            .id();
+
+        command_queue.apply(&mut world);
+        assert!(world.entities().len() == 1);
+        let results = world
+            .query::<(&u32, &u64)>()
+            .iter(&world)
+            .map(|(a, b)| (*a, *b))
+            .collect::<Vec<_>>();
+        assert_eq!(results, vec![(1u32, 4u64)]);
+    }
+
     #[test]
     fn remove_components() {
         let mut world = World::default();

crates/bevy_ecs/src/world/entity_ref.rs

@@ -194,76 +194,49 @@ impl<'w> EntityMut<'w> {
         let bundle_info = self.world.bundles.init_info::<T>(components);
         let current_location = self.location;
 
-        // Use a non-generic function to cut down on monomorphization
-        unsafe fn get_insert_bundle_info<'a>(
-            entities: &mut Entities,
-            archetypes: &'a mut Archetypes,
-            components: &mut Components,
-            storages: &mut Storages,
-            bundle_info: &BundleInfo,
-            current_location: EntityLocation,
-            entity: Entity,
-        ) -> (&'a Archetype, &'a Vec<ComponentStatus>, EntityLocation) {
-            // SAFE: component ids in `bundle_info` and self.location are valid
-            let new_archetype_id = add_bundle_to_archetype(
+        let (archetype, bundle_status, new_location) = unsafe {
+            Self::get_insert_bundle_info(
+                entities,
                 archetypes,
-                storages,
                 components,
-                current_location.archetype_id,
+                storages,
                 bundle_info,
-            );
-            if new_archetype_id == current_location.archetype_id {
-                let archetype = &archetypes[current_location.archetype_id];
-                let edge = archetype.edges().get_add_bundle(bundle_info.id).unwrap();
-                (archetype, &edge.bundle_status, current_location)
-            } else {
-                let (old_table_row, old_table_id) = {
-                    let old_archetype = &mut archetypes[current_location.archetype_id];
-                    let result = old_archetype.swap_remove(current_location.index);
-                    if let Some(swapped_entity) = result.swapped_entity {
-                        entities.meta[swapped_entity.id as usize].location = current_location;
-                    }
-                    (result.table_row, old_archetype.table_id())
-                };
-
-                let new_table_id = archetypes[new_archetype_id].table_id();
-
-                let new_location = if old_table_id == new_table_id {
-                    archetypes[new_archetype_id].allocate(entity, old_table_row)
-                } else {
-                    let (old_table, new_table) =
-                        storages.tables.get_2_mut(old_table_id, new_table_id);
-                    // PERF: store "non bundle" components in edge, then just move those to avoid
-                    // redundant copies
-                    let move_result =
-                        old_table.move_to_superset_unchecked(old_table_row, new_table);
-
-                    let new_location =
-                        archetypes[new_archetype_id].allocate(entity, move_result.new_row);
-                    // if an entity was moved into this entity's table spot, update its table row
-                    if let Some(swapped_entity) = move_result.swapped_entity {
-                        let swapped_location = entities.get(swapped_entity).unwrap();
-                        archetypes[swapped_location.archetype_id]
-                            .set_entity_table_row(swapped_location.index, old_table_row);
-                    }
-                    new_location
-                };
-
-                entities.meta[entity.id as usize].location = new_location;
-                let (old_archetype, new_archetype) =
-                    archetypes.get_2_mut(current_location.archetype_id, new_archetype_id);
-                let edge = old_archetype
-                    .edges()
-                    .get_add_bundle(bundle_info.id)
-                    .unwrap();
-                (&*new_archetype, &edge.bundle_status, new_location)
-
-                // Sparse set components are intentionally ignored here. They don't need to move
-            }
-        }
+                current_location,
+                entity,
+            )
+        };
+        self.location = new_location;
+
+        let table = &storages.tables[archetype.table_id()];
+        let table_row = archetype.entity_table_row(new_location.index);
+        // SAFE: table row is valid
+        unsafe {
+            bundle_info.write_components(
+                &mut storages.sparse_sets,
+                entity,
+                table,
+                table_row,
+                bundle_status,
+                bundle,
+                change_tick,
+            )
+        };
+        self
+    }
+
+    pub fn try_insert_bundle<T: Bundle>(&mut self, bundle: T) -> &mut Self {
+        let entity = self.entity;
+        let change_tick = self.world.change_tick();
+        let entities = &mut self.world.entities;
+        let archetypes = &mut self.world.archetypes;
+        let components = &mut self.world.components;
+        let storages = &mut self.world.storages;
+
+        let bundle_info = self.world.bundles.init_info::<T>(components);
+        let current_location = self.location;
 
         let (archetype, bundle_status, new_location) = unsafe {
-            get_insert_bundle_info(
+            Self::get_insert_bundle_info(
                 entities,
                 archetypes,
                 components,
@@ -279,7 +252,7 @@ impl<'w> EntityMut<'w> {
         let table_row = archetype.entity_table_row(new_location.index);
         // SAFE: table row is valid
         unsafe {
-            bundle_info.write_components(
+            bundle_info.write_additional_components(
                 &mut storages.sparse_sets,
                 entity,
                 table,
@@ -292,6 +265,72 @@ impl<'w> EntityMut<'w> {
         self
     }
 
+    // Use a non-generic function to cut down on monomorphization
+    unsafe fn get_insert_bundle_info<'a>(
+        entities: &mut Entities,
+        archetypes: &'a mut Archetypes,
+        components: &mut Components,
+        storages: &mut Storages,
+        bundle_info: &BundleInfo,
+        current_location: EntityLocation,
+        entity: Entity,
+    ) -> (&'a Archetype, &'a Vec<ComponentStatus>, EntityLocation) {
+        // SAFE: component ids in `bundle_info` and self.location are valid
+        let new_archetype_id = add_bundle_to_archetype(
+            archetypes,
+            storages,
+            components,
+            current_location.archetype_id,
+            bundle_info,
+        );
+        if new_archetype_id == current_location.archetype_id {
+            let archetype = &archetypes[current_location.archetype_id];
+            let edge = archetype.edges().get_add_bundle(bundle_info.id).unwrap();
+            (archetype, &edge.bundle_status, current_location)
+        } else {
+            let (old_table_row, old_table_id) = {
+                let old_archetype = &mut archetypes[current_location.archetype_id];
+                let result = old_archetype.swap_remove(current_location.index);
+                if let Some(swapped_entity) = result.swapped_entity {
+                    entities.meta[swapped_entity.id as usize].location = current_location;
+                }
+                (result.table_row, old_archetype.table_id())
+            };
+
+            let new_table_id = archetypes[new_archetype_id].table_id();
+
+            let new_location = if old_table_id == new_table_id {
+                archetypes[new_archetype_id].allocate(entity, old_table_row)
+            } else {
+                let (old_table, new_table) = storages.tables.get_2_mut(old_table_id, new_table_id);
+                // PERF: store "non bundle" components in edge, then just move those to avoid
+                // redundant copies
+                let move_result = old_table.move_to_superset_unchecked(old_table_row, new_table);
+
+                let new_location =
+                    archetypes[new_archetype_id].allocate(entity, move_result.new_row);
+                // if an entity was moved into this entity's table spot, update its table row
+                if let Some(swapped_entity) = move_result.swapped_entity {
+                    let swapped_location = entities.get(swapped_entity).unwrap();
+                    archetypes[swapped_location.archetype_id]
+                        .set_entity_table_row(swapped_location.index, old_table_row);
+                }
+                new_location
+            };
+
+            entities.meta[entity.id as usize].location = new_location;
+            let (old_archetype, new_archetype) =
+                archetypes.get_2_mut(current_location.archetype_id, new_archetype_id);
+            let edge = old_archetype
+                .edges()
+                .get_add_bundle(bundle_info.id)
+                .unwrap();
+            (&*new_archetype, &edge.bundle_status, new_location)
+
+            // Sparse set components are intentionally ignored here. They don't need to move
+        }
+    }
+
     pub fn remove_bundle<T: Bundle>(&mut self) -> Option<T> {
         let archetypes = &mut self.world.archetypes;
         let storages = &mut self.world.storages;
@@ -461,6 +500,10 @@ impl<'w> EntityMut<'w> {
         self.insert_bundle((value,))
     }
 
+    pub fn try_insert<T: Component>(&mut self, value: T) -> &mut Self {
+        self.try_insert_bundle((value,))
+    }
+
     pub fn remove<T: Component>(&mut self) -> Option<T> {
         self.remove_bundle::<(T,)>().map(|v| v.0)
     }