From caa662272ce91028cef6e237bb7c2a3e78727ef4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20=C5=81abor?= Date: Mon, 20 Mar 2023 20:33:11 -0400 Subject: [PATCH] bevy_ecs: add untyped methods for inserting components and bundles (#7204) This MR is a rebased and alternative proposal to https://github.com/bevyengine/bevy/pull/5602 # Objective - https://github.com/bevyengine/bevy/pull/4447 implemented untyped (using component ids instead of generics and TypeId) APIs for inserting/accessing resources and accessing components, but left inserting components for another PR (this one) ## Solution - add `EntityMut::insert_by_id` - split `Bundle` into `DynamicBundle` with `get_components` and `Bundle: DynamicBundle`. This allows the `BundleInserter` machinery to be reused for bundles that can only be written, not read, and have no statically available `ComponentIds` - Compared to the original MR this approach exposes unsafe endpoints and requires the user to manage instantiated `BundleIds`. This is quite easy for the end user to do and does not incur the performance penalty of checking whether component input is correctly provided for the `BundleId`. - This MR does ensure that constructing `BundleId` itself is safe --- ## Changelog - add methods for inserting bundles and components to: `world.entity_mut(entity).insert_by_id` --- crates/bevy_ecs/macros/src/lib.rs | 2 + crates/bevy_ecs/src/bundle.rs | 110 ++++++++++++- crates/bevy_ecs/src/world/entity_ref.rs | 201 +++++++++++++++++++++++- 3 files changed, 303 insertions(+), 10 deletions(-) diff --git a/crates/bevy_ecs/macros/src/lib.rs b/crates/bevy_ecs/macros/src/lib.rs index 0e3374f757c83..a59c54db8d1ef 100644 --- a/crates/bevy_ecs/macros/src/lib.rs +++ b/crates/bevy_ecs/macros/src/lib.rs @@ -126,7 +126,9 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream { #(#field_from_components)* } } + } + impl #impl_generics #ecs_path::bundle::DynamicBundle for #struct_name #ty_generics #where_clause { #[allow(unused_variables)] #[inline] fn get_components( diff --git a/crates/bevy_ecs/src/bundle.rs b/crates/bevy_ecs/src/bundle.rs index ae37fdcdccd0e..9735c25271c98 100644 --- a/crates/bevy_ecs/src/bundle.rs +++ b/crates/bevy_ecs/src/bundle.rs @@ -3,7 +3,7 @@ //! This module contains the [`Bundle`] trait and some other helper types. pub use bevy_ecs_macros::Bundle; -use bevy_utils::HashSet; +use bevy_utils::{HashMap, HashSet}; use crate::{ archetype::{ @@ -135,10 +135,10 @@ use std::any::TypeId; /// [`Query`]: crate::system::Query // Some safety points: // - [`Bundle::component_ids`] must return the [`ComponentId`] for each component type in the -// bundle, in the _exact_ order that [`Bundle::get_components`] is called. +// bundle, in the _exact_ order that [`DynamicBundle::get_components`] is called. // - [`Bundle::from_components`] must call `func` exactly once for each [`ComponentId`] returned by // [`Bundle::component_ids`]. -pub unsafe trait Bundle: Send + Sync + 'static { +pub unsafe trait Bundle: DynamicBundle + Send + Sync + 'static { /// Gets this [`Bundle`]'s component ids, in the order of this bundle's [`Component`]s #[doc(hidden)] fn component_ids( @@ -159,7 +159,10 @@ pub unsafe trait Bundle: Send + Sync + 'static { // Ensure that the `OwningPtr` is used correctly F: for<'a> FnMut(&'a mut T) -> OwningPtr<'a>, Self: Sized; +} +/// The parts from [`Bundle`] that don't require statically knowing the components of the bundle. +pub trait DynamicBundle { // SAFETY: // The `StorageType` argument passed into [`Bundle::get_components`] must be correct for the // component being fetched. @@ -192,7 +195,9 @@ unsafe impl Bundle for C { // Safety: The id given in `component_ids` is for `Self` func(ctx).read() } +} +impl DynamicBundle for C { #[inline] fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) { OwningPtr::make(self, |ptr| func(C::Storage::STORAGE_TYPE, ptr)); @@ -203,7 +208,7 @@ macro_rules! tuple_impl { ($($name: ident),*) => { // SAFETY: // - `Bundle::component_ids` calls `ids` for each component type in the - // bundle, in the exact order that `Bundle::get_components` is called. + // bundle, in the exact order that `DynamicBundle::get_components` is called. // - `Bundle::from_components` calls `func` exactly once for each `ComponentId` returned by `Bundle::component_ids`. // - `Bundle::get_components` is called exactly once for each member. Relies on the above implementation to pass the correct // `StorageType` into the callback. @@ -223,7 +228,9 @@ macro_rules! tuple_impl { // https://doc.rust-lang.org/reference/expressions.html#evaluation-order-of-operands ($(<$name as Bundle>::from_components(ctx, func),)*) } + } + impl<$($name: Bundle),*> DynamicBundle for ($($name,)*) { #[allow(unused_variables, unused_mut)] #[inline(always)] fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) { @@ -376,7 +383,7 @@ impl BundleInfo { /// `entity`, `bundle` must match this [`BundleInfo`]'s type #[inline] #[allow(clippy::too_many_arguments)] - unsafe fn write_components( + unsafe fn write_components( &self, table: &mut Table, sparse_sets: &mut SparseSets, @@ -527,7 +534,7 @@ impl<'a, 'b> BundleInserter<'a, 'b> { /// `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( + pub unsafe fn insert( &mut self, entity: Entity, location: EntityLocation, @@ -677,7 +684,7 @@ impl<'a, 'b> BundleSpawner<'a, 'b> { /// # Safety /// `entity` must be allocated (but non-existent), `T` must match this [`BundleInfo`]'s type #[inline] - pub unsafe fn spawn_non_existent( + pub unsafe fn spawn_non_existent( &mut self, entity: Entity, bundle: T, @@ -712,7 +719,12 @@ impl<'a, 'b> BundleSpawner<'a, 'b> { #[derive(Default)] pub struct Bundles { bundle_infos: Vec, + /// Cache static [`BundleId`] bundle_ids: TypeIdMap, + /// Cache dynamic [`BundleId`] with multiple components + dynamic_bundle_ids: HashMap, (BundleId, Vec)>, + /// Cache optimized dynamic [`BundleId`] with single component + dynamic_component_bundle_ids: HashMap, } impl Bundles { @@ -726,6 +738,7 @@ impl Bundles { self.bundle_ids.get(&type_id).cloned() } + /// Initializes a new [`BundleInfo`] for a statically known type. pub(crate) fn init_info<'a, T: Bundle>( &'a mut self, components: &mut Components, @@ -745,6 +758,64 @@ impl Bundles { // SAFETY: index either exists, or was initialized unsafe { self.bundle_infos.get_unchecked(id.0) } } + + /// Initializes a new [`BundleInfo`] for a dynamic [`Bundle`]. + /// + /// # Panics + /// + /// Panics if any of the provided [`ComponentId`]s do not exist in the + /// provided [`Components`]. + pub(crate) fn init_dynamic_info( + &mut self, + components: &mut Components, + component_ids: &[ComponentId], + ) -> (&BundleInfo, &Vec) { + let bundle_infos = &mut self.bundle_infos; + + // Use `raw_entry_mut` to avoid cloning `component_ids` to access `Entry` + let (_, (bundle_id, storage_types)) = self + .dynamic_bundle_ids + .raw_entry_mut() + .from_key(component_ids) + .or_insert_with(|| { + ( + Vec::from(component_ids), + initialize_dynamic_bundle(bundle_infos, components, Vec::from(component_ids)), + ) + }); + + // SAFETY: index either exists, or was initialized + let bundle_info = unsafe { bundle_infos.get_unchecked(bundle_id.0) }; + + (bundle_info, storage_types) + } + + /// Initializes a new [`BundleInfo`] for a dynamic [`Bundle`] with single component. + /// + /// # Panics + /// + /// Panics if the provided [`ComponentId`] does not exist in the provided [`Components`]. + pub(crate) fn init_component_info( + &mut self, + components: &mut Components, + component_id: ComponentId, + ) -> (&BundleInfo, StorageType) { + let bundle_infos = &mut self.bundle_infos; + let (bundle_id, storage_types) = self + .dynamic_component_bundle_ids + .entry(component_id) + .or_insert_with(|| { + let (id, storage_type) = + initialize_dynamic_bundle(bundle_infos, components, vec![component_id]); + // SAFETY: `storage_type` guaranteed to have length 1 + (id, storage_type[0]) + }); + + // SAFETY: index either exists, or was initialized + let bundle_info = unsafe { bundle_infos.get_unchecked(bundle_id.0) }; + + (bundle_info, *storage_types) + } } /// # Safety @@ -784,3 +855,28 @@ unsafe fn initialize_bundle( BundleInfo { id, component_ids } } + +/// Asserts that all components are part of [`Components`] +/// and initializes a [`BundleInfo`]. +fn initialize_dynamic_bundle( + bundle_infos: &mut Vec, + components: &Components, + component_ids: Vec, +) -> (BundleId, Vec) { + // Assert component existence + let storage_types = component_ids.iter().map(|&id| { + components.get_info(id).unwrap_or_else(|| { + panic!( + "init_dynamic_info called with component id {id:?} which doesn't exist in this world" + ) + }).storage_type() + }).collect(); + + let id = BundleId(bundle_infos.len()); + let bundle_info = + // SAFETY: `component_ids` are valid as they were just checked + unsafe { initialize_bundle("", components, component_ids, id) }; + bundle_infos.push(bundle_info); + + (id, storage_types) +} diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 5a60d0a7b5721..39e5b8e40f63b 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -1,6 +1,6 @@ use crate::{ archetype::{Archetype, ArchetypeId, Archetypes}, - bundle::{Bundle, BundleInfo}, + bundle::{Bundle, BundleInfo, BundleInserter, DynamicBundle}, change_detection::MutUntyped, component::{Component, ComponentId, ComponentTicks, Components, StorageType}, entity::{Entities, Entity, EntityLocation}, @@ -279,6 +279,90 @@ impl<'w> EntityMut<'w> { self } + /// Inserts a dynamic [`Component`] into the entity. + /// + /// This will overwrite any previous value(s) of the same component type. + /// + /// You should prefer to use the typed API [`EntityMut::insert`] where possible. + /// + /// # Safety + /// + /// - [`ComponentId`] must be from the same world as [`EntityMut`] + /// - [`OwningPtr`] must be a valid reference to the type represented by [`ComponentId`] + pub unsafe fn insert_by_id( + &mut self, + component_id: ComponentId, + component: OwningPtr<'_>, + ) -> &mut Self { + let change_tick = self.world.change_tick(); + + let bundles = &mut self.world.bundles; + let components = &mut self.world.components; + + let (bundle_info, storage_type) = bundles.init_component_info(components, component_id); + let 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, + ); + + self.location = insert_dynamic_bundle( + bundle_inserter, + self.entity, + self.location, + Some(component).into_iter(), + Some(storage_type).into_iter(), + ); + + self + } + + /// Inserts a dynamic [`Bundle`] into the entity. + /// + /// This will overwrite any previous value(s) of the same component type. + /// + /// You should prefer to use the typed API [`EntityMut::insert`] where possible. + /// If your [`Bundle`] only has one component, use the cached API [`EntityMut::insert_by_id`]. + /// + /// If possible, pass a sorted slice of `ComponentId` to maximize caching potential. + /// + /// # Safety + /// - Each [`ComponentId`] must be from the same world as [`EntityMut`] + /// - Each [`OwningPtr`] must be a valid reference to the type represented by [`ComponentId`] + pub unsafe fn insert_by_ids<'a, I: Iterator>>( + &mut self, + component_ids: &[ComponentId], + iter_components: I, + ) -> &mut Self { + let change_tick = self.world.change_tick(); + + let bundles = &mut self.world.bundles; + let components = &mut self.world.components; + + let (bundle_info, storage_types) = bundles.init_dynamic_info(components, component_ids); + let 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, + ); + + self.location = insert_dynamic_bundle( + bundle_inserter, + self.entity, + self.location, + iter_components, + storage_types.iter().cloned(), + ); + + self + } + /// Removes all components in the [`Bundle`] from the entity and returns their previous values. /// /// **Note:** If the entity does not have every component in the bundle, this method will not @@ -603,7 +687,7 @@ impl<'w> EntityMut<'w> { /// // Mutate the world while we have access to it. /// let mut r = world.resource_mut::(); /// r.0 += 1; - /// + /// /// // Return a value from the world before giving it back to the `EntityMut`. /// *r /// }); @@ -672,6 +756,44 @@ impl<'w> EntityMut<'w> { } } +/// Inserts a dynamic [`Bundle`] into the entity. +/// +/// # Safety +/// +/// - [`OwningPtr`] and [`StorageType`] iterators must correspond to the +/// [`BundleInfo`] used to construct [`BundleInserter`] +/// - [`Entity`] must correspond to [`EntityLocation`] +unsafe fn insert_dynamic_bundle< + 'a, + I: Iterator>, + S: Iterator, +>( + mut bundle_inserter: BundleInserter<'_, '_>, + entity: Entity, + location: EntityLocation, + components: I, + storage_types: S, +) -> EntityLocation { + struct DynamicInsertBundle<'a, I: Iterator)>> { + components: I, + } + + impl<'a, I: Iterator)>> DynamicBundle + for DynamicInsertBundle<'a, I> + { + fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) { + self.components.for_each(|(t, ptr)| func(t, ptr)); + } + } + + let bundle = DynamicInsertBundle { + components: storage_types.zip(components), + }; + + // SAFETY: location matches current entity. + unsafe { bundle_inserter.insert(entity, location, bundle) } +} + /// Removes a bundle from the given archetype and returns the resulting archetype (or None if the /// removal was invalid). 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. @@ -835,6 +957,7 @@ pub(crate) unsafe fn take_component<'a>( #[cfg(test)] mod tests { + use bevy_ptr::OwningPtr; use std::panic::AssertUnwindSafe; use crate as bevy_ecs; @@ -862,9 +985,13 @@ mod tests { assert_eq!(a, vec![1]); } - #[derive(Component)] + #[derive(Component, Clone, Copy, Debug, PartialEq)] struct TestComponent(u32); + #[derive(Component, Clone, Copy, Debug, PartialEq)] + #[component(storage = "SparseSet")] + struct TestComponent2(u32); + #[test] fn entity_ref_get_by_id() { let mut world = World::new(); @@ -1107,4 +1234,72 @@ mod tests { assert_eq!(world.entity(e2).get::().unwrap(), &Dense(1)); } + + #[test] + fn entity_mut_insert_by_id() { + let mut world = World::new(); + let test_component_id = world.init_component::(); + + let mut entity = world.spawn_empty(); + OwningPtr::make(TestComponent(42), |ptr| { + // SAFETY: `ptr` matches the component id + unsafe { entity.insert_by_id(test_component_id, ptr) }; + }); + + let components: Vec<_> = world.query::<&TestComponent>().iter(&world).collect(); + + assert_eq!(components, vec![&TestComponent(42)]); + + // Compare with `insert_bundle_by_id` + + let mut entity = world.spawn_empty(); + OwningPtr::make(TestComponent(84), |ptr| { + // SAFETY: `ptr` matches the component id + unsafe { entity.insert_by_ids(&[test_component_id], vec![ptr].into_iter()) }; + }); + + let components: Vec<_> = world.query::<&TestComponent>().iter(&world).collect(); + + assert_eq!(components, vec![&TestComponent(42), &TestComponent(84)]); + } + + #[test] + fn entity_mut_insert_bundle_by_id() { + let mut world = World::new(); + let test_component_id = world.init_component::(); + let test_component_2_id = world.init_component::(); + + let component_ids = [test_component_id, test_component_2_id]; + let test_component_value = TestComponent(42); + let test_component_2_value = TestComponent2(84); + + let mut entity = world.spawn_empty(); + OwningPtr::make(test_component_value, |ptr1| { + OwningPtr::make(test_component_2_value, |ptr2| { + // SAFETY: `ptr1` and `ptr2` match the component ids + unsafe { entity.insert_by_ids(&component_ids, vec![ptr1, ptr2].into_iter()) }; + }); + }); + + let dynamic_components: Vec<_> = world + .query::<(&TestComponent, &TestComponent2)>() + .iter(&world) + .collect(); + + assert_eq!( + dynamic_components, + vec![(&TestComponent(42), &TestComponent2(84))] + ); + + // Compare with `World` generated using static type equivalents + let mut static_world = World::new(); + + static_world.spawn((test_component_value, test_component_2_value)); + let static_components: Vec<_> = static_world + .query::<(&TestComponent, &TestComponent2)>() + .iter(&static_world) + .collect(); + + assert_eq!(dynamic_components, static_components); + } }