From e8939755177f15ce0c1d067b3ee0f4a8cd43d16c Mon Sep 17 00:00:00 2001 From: jfaz Date: Fri, 21 Jul 2023 12:10:17 -0400 Subject: [PATCH 1/2] Add `insert_unique` command --- crates/bevy_ecs/src/bundle.rs | 207 +++++++++++++++++++++ crates/bevy_ecs/src/storage/sparse_set.rs | 24 +++ crates/bevy_ecs/src/system/commands/mod.rs | 78 ++++++++ crates/bevy_ecs/src/world/entity_ref.rs | 46 +++++ 4 files changed, 355 insertions(+) diff --git a/crates/bevy_ecs/src/bundle.rs b/crates/bevy_ecs/src/bundle.rs index f121e3429c5d9..598648ac98773 100644 --- a/crates/bevy_ecs/src/bundle.rs +++ b/crates/bevy_ecs/src/bundle.rs @@ -488,6 +488,65 @@ impl BundleInfo { }); } + /// This writes components from a given [`Bundle`] to the given entity, without overwriting existing component types. + /// + /// # Safety + /// + /// `bundle_component_status` must return the "correct" [`ComponentStatus`] for each component + /// in the [`Bundle`], with respect to the entity's original archetype (prior to the bundle being added) + /// For example, if the original archetype already has `ComponentA` and `T` also has `ComponentA`, the status + /// should be `Mutated`. If the original archetype does not have `ComponentA`, the status should be `Added`. + /// When "inserting" a bundle into an existing entity, [`AddBundle`](crate::archetype::AddBundle) + /// should be used, which will report `Added` vs `Mutated` status based on the current archetype's structure. + /// When spawning a bundle, [`SpawnBundleStatus`] can be used instead, which removes the need + /// to look up the [`AddBundle`](crate::archetype::AddBundle) in the archetype graph, which requires + /// ownership of the entity's current archetype. + /// + /// `table` must be the "new" table for `entity`. `table_row` must have space allocated for the + /// `entity`, `bundle` must match this [`BundleInfo`]'s type + #[inline] + #[allow(clippy::too_many_arguments)] + unsafe fn write_components_unique( + &self, + table: &mut Table, + sparse_sets: &mut SparseSets, + bundle_component_status: &S, + entity: Entity, + table_row: TableRow, + change_tick: Tick, + bundle: T, + ) { + // 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(&mut |storage_type, component_ptr| { + let component_id = *self.component_ids.get_unchecked(bundle_component); + match storage_type { + StorageType::Table => { + let column = + // SAFETY: If component_id is in self.component_ids, BundleInfo::new requires that + // the target table contains the component. + unsafe { table.get_column_mut(component_id).debug_checked_unwrap() }; + // SAFETY: bundle_component is a valid index for this bundle + match bundle_component_status.get_status(bundle_component) { + ComponentStatus::Added => { + column.initialize(table_row, component_ptr, change_tick); + } + _ => (), + } + } + StorageType::SparseSet => { + let sparse_set = + // SAFETY: If component_id is in self.component_ids, BundleInfo::new requires that + // a sparse set exists for the component. + unsafe { sparse_sets.get_mut(component_id).debug_checked_unwrap() }; + sparse_set.insert_unique(entity, component_ptr, change_tick); + } + } + bundle_component += 1; + }); + } + /// Adds a bundle to the given archetype and returns the resulting archetype. This could be the /// same [`ArchetypeId`], in the event that adding the given bundle does not result in an /// [`Archetype`] change. Results are cached in the [`Archetype`] graph to avoid redundant work. @@ -742,6 +801,154 @@ impl<'a, 'b> BundleInserter<'a, 'b> { } } } + + /// # Safety + /// `entity` must currently exist in the source archetype for this inserter. `archetype_row` + /// must be `entity`'s location in the archetype. `T` must match this [`BundleInfo`]'s type + #[inline] + pub unsafe fn insert_unique( + &mut self, + entity: Entity, + location: EntityLocation, + bundle: T, + ) -> EntityLocation { + match &mut self.result { + InsertBundleResult::SameArchetype => { + // PERF: this could be looked up during Inserter construction and stored (but borrowing makes this nasty) + // SAFETY: The edge is assured to be initialized when creating the BundleInserter + let add_bundle = unsafe { + self.archetype + .edges() + .get_add_bundle_internal(self.bundle_info.id) + .debug_checked_unwrap() + }; + self.bundle_info.write_components_unique( + self.table, + self.sparse_sets, + add_bundle, + entity, + location.table_row, + self.change_tick, + bundle, + ); + location + } + InsertBundleResult::NewArchetypeSameTable { new_archetype } => { + let result = self.archetype.swap_remove(location.archetype_row); + if let Some(swapped_entity) = result.swapped_entity { + let swapped_location = + // SAFETY: If the swap was successful, swapped_entity must be valid. + unsafe { self.entities.get(swapped_entity).debug_checked_unwrap() }; + self.entities.set( + swapped_entity.index(), + EntityLocation { + archetype_id: swapped_location.archetype_id, + archetype_row: location.archetype_row, + table_id: swapped_location.table_id, + table_row: swapped_location.table_row, + }, + ); + } + let new_location = new_archetype.allocate(entity, result.table_row); + self.entities.set(entity.index(), new_location); + + // PERF: this could be looked up during Inserter construction and stored (but borrowing makes this nasty) + // SAFETY: The edge is assured to be initialized when creating the BundleInserter + let add_bundle = unsafe { + self.archetype + .edges() + .get_add_bundle_internal(self.bundle_info.id) + .debug_checked_unwrap() + }; + self.bundle_info.write_components_unique( + self.table, + self.sparse_sets, + add_bundle, + entity, + result.table_row, + self.change_tick, + bundle, + ); + new_location + } + InsertBundleResult::NewArchetypeNewTable { + new_archetype, + new_table, + } => { + let result = self.archetype.swap_remove(location.archetype_row); + if let Some(swapped_entity) = result.swapped_entity { + let swapped_location = + // SAFETY: If the swap was successful, swapped_entity must be valid. + unsafe { self.entities.get(swapped_entity).debug_checked_unwrap() }; + self.entities.set( + swapped_entity.index(), + EntityLocation { + archetype_id: swapped_location.archetype_id, + archetype_row: location.archetype_row, + table_id: swapped_location.table_id, + table_row: swapped_location.table_row, + }, + ); + } + // PERF: store "non bundle" components in edge, then just move those to avoid + // redundant copies + let move_result = self + .table + .move_to_superset_unchecked(result.table_row, new_table); + let new_location = new_archetype.allocate(entity, move_result.new_row); + self.entities.set(entity.index(), new_location); + + // 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 = + // SAFETY: If the swap was successful, swapped_entity must be valid. + unsafe { self.entities.get(swapped_entity).debug_checked_unwrap() }; + let swapped_archetype = if self.archetype.id() == swapped_location.archetype_id + { + &mut *self.archetype + } else if new_archetype.id() == swapped_location.archetype_id { + new_archetype + } else { + // SAFETY: the only two borrowed archetypes are above and we just did collision checks + &mut *self + .archetypes_ptr + .add(swapped_location.archetype_id.index()) + }; + + self.entities.set( + swapped_entity.index(), + EntityLocation { + archetype_id: swapped_location.archetype_id, + archetype_row: swapped_location.archetype_row, + table_id: swapped_location.table_id, + table_row: result.table_row, + }, + ); + swapped_archetype + .set_entity_table_row(swapped_location.archetype_row, result.table_row); + } + + // PERF: this could be looked up during Inserter construction and stored (but borrowing makes this nasty) + // SAFETY: The edge is assured to be initialized when creating the BundleInserter + let add_bundle = unsafe { + self.archetype + .edges() + .get_add_bundle_internal(self.bundle_info.id) + .debug_checked_unwrap() + }; + self.bundle_info.write_components_unique( + new_table, + self.sparse_sets, + add_bundle, + entity, + move_result.new_row, + self.change_tick, + bundle, + ); + new_location + } + } + } } pub(crate) struct BundleSpawner<'a, 'b> { diff --git a/crates/bevy_ecs/src/storage/sparse_set.rs b/crates/bevy_ecs/src/storage/sparse_set.rs index 4d318d1a4677f..758db33e32b4c 100644 --- a/crates/bevy_ecs/src/storage/sparse_set.rs +++ b/crates/bevy_ecs/src/storage/sparse_set.rs @@ -187,6 +187,30 @@ impl ComponentSparseSet { } } + /// Inserts the `entity` key and component `value` pair into this sparse set if they do not exist + /// + /// # Safety + /// The `value` pointer must point to a valid address that matches the [`Layout`](std::alloc::Layout) + /// inside the [`ComponentInfo`] given when constructing this sparse set. + pub(crate) unsafe fn insert_unique( + &mut self, + entity: Entity, + value: OwningPtr<'_>, + change_tick: Tick, + ) { + if self.sparse.get(entity.index()).is_none() { + let dense_index = self.dense.len(); + self.dense.push(value, ComponentTicks::new(change_tick)); + self.sparse.insert(entity.index(), dense_index as u32); + #[cfg(debug_assertions)] + assert_eq!(self.entities.len(), dense_index); + #[cfg(not(debug_assertions))] + self.entities.push(entity.index()); + #[cfg(debug_assertions)] + self.entities.push(entity); + } + } + /// Returns `true` if the sparse set has a component value for the provided `entity`. #[inline] pub fn contains(&self, entity: Entity) -> bool { diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs index fac21c7c0ffc2..bfbec2e912fba 100644 --- a/crates/bevy_ecs/src/system/commands/mod.rs +++ b/crates/bevy_ecs/src/system/commands/mod.rs @@ -722,6 +722,63 @@ impl<'w, 's, 'a> EntityCommands<'w, 's, 'a> { self } + /// Adds a [`Bundle`] of components to the entity. + /// + /// This will preserve any previous value(s) of the same component type. + /// + /// # Panics + /// + /// The command will panic when applied if the associated entity does not exist. + /// + /// # Example + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # #[derive(Resource)] + /// # struct PlayerEntity { entity: Entity } + /// #[derive(Component)] + /// struct Health(u32); + /// #[derive(Component)] + /// struct Strength(u32); + /// #[derive(Component)] + /// struct Defense(u32); + /// + /// #[derive(Bundle)] + /// struct CombatBundle { + /// health: Health, + /// strength: Strength, + /// } + /// + /// fn add_combat_stats_system(mut commands: Commands, player: Res) { + /// commands + /// .entity(player.entity) + /// // You can insert individual components: + /// .insert_unique(Defense(10)) + /// // You can also insert pre-defined bundles of components: + /// .insert_unique(CombatBundle { + /// health: Health(100), + /// strength: Strength(40), + /// }) + /// // You can also insert tuples of components and bundles. + /// // This is equivalent to the calls above: + /// .insert_unique(( + /// Defense(10), + /// CombatBundle { + /// health: Health(100), + /// strength: Strength(40), + /// }, + /// )); + /// } + /// # bevy_ecs::system::assert_is_system(add_combat_stats_system); + /// ``` + pub fn insert_unique(&mut self, bundle: impl Bundle) -> &mut Self { + self.commands.add(InsertUnique { + entity: self.entity, + bundle, + }); + self + } + /// Removes a [`Bundle`] of components from the entity. /// /// See [`EntityMut::remove`](crate::world::EntityMut::remove) for more @@ -959,6 +1016,27 @@ where } } +/// A [`Command`] that adds the components in a [`Bundle`] to an entity, without overwriting any existing components on the original entity +pub struct InsertUnique { + /// The entity to which the components will be added. + pub entity: Entity, + /// The [`Bundle`] containing the components that will be added to the entity. + pub bundle: T, +} + +impl Command for InsertUnique +where + T: Bundle + 'static, +{ + fn apply(self, world: &mut World) { + if let Some(mut entity) = world.get_entity_mut(self.entity) { + entity.insert_unique(self.bundle); + } else { + panic!("error[B0003]: Could not insert a bundle (of type `{}`) for entity {:?} because it doesn't exist in this World.", std::any::type_name::(), self.entity); + } + } +} + /// A [`Command`] that removes components from an entity. /// For a [`Bundle`] type `T`, this will remove any components in the bundle. /// Any components in the bundle that aren't found on the entity will be ignored. diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 29a63c2220d60..bc28569963845 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -346,6 +346,31 @@ impl<'w> EntityMut<'w> { self } + /// Adds a [`Bundle`] of components to the entity. + /// + /// This will preserve any previous value(s) of the same component type. + pub fn insert_unique(&mut self, bundle: T) -> &mut Self { + let change_tick = self.world.change_tick(); + let bundle_info = self + .world + .bundles + .init_info::(&mut self.world.components, &mut self.world.storages); + let mut bundle_inserter = bundle_info.get_bundle_inserter( + &mut self.world.entities, + &mut self.world.archetypes, + &mut self.world.components, + &mut self.world.storages, + self.location.archetype_id, + change_tick, + ); + // SAFETY: location matches current entity. `T` matches `bundle_info` + unsafe { + self.location = bundle_inserter.insert_unique(self.entity, self.location, bundle); + } + + self + } + /// Inserts a dynamic [`Component`] into the entity. /// /// This will overwrite any previous value(s) of the same component type. @@ -1369,4 +1394,25 @@ mod tests { assert_eq!(dynamic_components, static_components); } + + #[test] + fn entity_mut_insert_bundle_unique() { + #[derive(Component)] + struct TableComponent(u8); + #[derive(Component)] + #[component(storage = "SparseSet")] + struct SparseComponent(u8); + + let mut world = World::new(); + let mut test_entity = world.spawn_empty(); + + // First time inserting a component type should succeed, any further inserts of the same type should be dropped + test_entity.insert_unique(TableComponent(7)); + test_entity.insert_unique((TableComponent(3), SparseComponent(7))); + test_entity.insert_unique((TableComponent(1), SparseComponent(1))); + test_entity.insert_unique((TableComponent(5), SparseComponent(5))); + + assert_eq!(test_entity.get::().unwrap().0, 7); + assert_eq!(test_entity.get::().unwrap().0, 7); + } } From 7a0cce9adb9270ce6b06ec6dad4974265934174c Mon Sep 17 00:00:00 2001 From: jfaz Date: Fri, 21 Jul 2023 13:42:11 -0400 Subject: [PATCH 2/2] Use `if let` for destructuring --- crates/bevy_ecs/src/bundle.rs | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/crates/bevy_ecs/src/bundle.rs b/crates/bevy_ecs/src/bundle.rs index 598648ac98773..290f2b76a0899 100644 --- a/crates/bevy_ecs/src/bundle.rs +++ b/crates/bevy_ecs/src/bundle.rs @@ -528,11 +528,10 @@ impl BundleInfo { // the target table contains the component. unsafe { table.get_column_mut(component_id).debug_checked_unwrap() }; // SAFETY: bundle_component is a valid index for this bundle - match bundle_component_status.get_status(bundle_component) { - ComponentStatus::Added => { - column.initialize(table_row, component_ptr, change_tick); - } - _ => (), + if let ComponentStatus::Added = + bundle_component_status.get_status(bundle_component) + { + column.initialize(table_row, component_ptr, change_tick); } } StorageType::SparseSet => {