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

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

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<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)
+}

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<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
@@ -603,7 +687,7 @@ impl<'w> EntityMut<'w> {
     ///     // Mutate the world while we have access to it.
     ///     let mut r = world.resource_mut::<R>();
     ///     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<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);
+    }
 }