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`
This commit is contained in:
Jakub Łabor 2023-03-20 20:33:11 -04:00 committed by GitHub
parent 3b51e1c8d9
commit caa662272c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 303 additions and 10 deletions

View File

@ -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(

View File

@ -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<C: Component> Bundle for C {
// Safety: The id given in `component_ids` is for `Self`
func(ctx).read()
}
}
impl<C: Component> 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<T: Bundle, S: BundleComponentStatus>(
unsafe fn write_components<T: DynamicBundle, S: BundleComponentStatus>(
&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<T: Bundle>(
pub unsafe fn insert<T: DynamicBundle>(
&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<T: Bundle>(
pub unsafe fn spawn_non_existent<T: DynamicBundle>(
&mut self,
entity: Entity,
bundle: T,
@ -712,7 +719,12 @@ impl<'a, 'b> BundleSpawner<'a, 'b> {
#[derive(Default)]
pub struct Bundles {
bundle_infos: Vec<BundleInfo>,
/// Cache static [`BundleId`]
bundle_ids: TypeIdMap<BundleId>,
/// Cache dynamic [`BundleId`] with multiple components
dynamic_bundle_ids: HashMap<Vec<ComponentId>, (BundleId, Vec<StorageType>)>,
/// Cache optimized dynamic [`BundleId`] with single component
dynamic_component_bundle_ids: HashMap<ComponentId, (BundleId, StorageType)>,
}
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<StorageType>) {
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<BundleInfo>,
components: &Components,
component_ids: Vec<ComponentId>,
) -> (BundleId, Vec<StorageType>) {
// 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("<dynamic bundle>", components, component_ids, id) };
bundle_infos.push(bundle_info);
(id, storage_types)
}

View File

@ -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<Item = OwningPtr<'a>>>(
&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
@ -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<Item = OwningPtr<'a>>,
S: Iterator<Item = StorageType>,
>(
mut bundle_inserter: BundleInserter<'_, '_>,
entity: Entity,
location: EntityLocation,
components: I,
storage_types: S,
) -> EntityLocation {
struct DynamicInsertBundle<'a, I: Iterator<Item = (StorageType, OwningPtr<'a>)>> {
components: I,
}
impl<'a, I: Iterator<Item = (StorageType, OwningPtr<'a>)>> 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::<Dense>().unwrap(), &Dense(1));
}
#[test]
fn entity_mut_insert_by_id() {
let mut world = World::new();
let test_component_id = world.init_component::<TestComponent>();
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::<TestComponent>();
let test_component_2_id = world.init_component::<TestComponent2>();
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);
}
}