Remove unnecessary branching from bundle insertion (#6902)

# Objective
Speed up bundle insertion and spawning from a bundle.

## Solution
Use the same technique used in #6800 to remove the branch on storage type when writing components from a `Bundle` into storage.

 - Add a `StorageType` argument to the closure on `Bundle::get_components`.
 - Pass `C::Storage::STORAGE_TYPE` into that argument.
 - Match on that argument instead of reading from a `Vec<StorageType>` in `BundleInfo`.
 - Marked all implementations of `Bundle::get_components` as inline to encourage dead code elimination.

The `Vec<StorageType>` in `BundleInfo` was also removed as it's no longer needed. If users were reliant on this, they can either use the compile time constants or fetch the information from `Components`. Should save a rather negligible amount of memory.

## Performance
Microbenchmarks show a slight improvement to inserting components into existing entities, as well as spawning from a bundle. Ranging about 8-16% faster depending on the benchmark.

```
group                                          main                                    soft-constant-write-components
-----                                          ----                                    ------------------------------
add_remove/sparse_set                          1.08  1019.0±80.10µs        ? ?/sec     1.00   944.6±66.86µs        ? ?/sec
add_remove/table                               1.07  1343.3±20.37µs        ? ?/sec     1.00  1257.3±18.13µs        ? ?/sec
add_remove_big/sparse_set                      1.08  1132.4±263.10µs        ? ?/sec    1.00  1050.8±240.74µs        ? ?/sec
add_remove_big/table                           1.02      2.6±0.05ms        ? ?/sec     1.00      2.5±0.08ms        ? ?/sec
get_or_spawn/batched                           1.15   401.4±17.76µs        ? ?/sec     1.00   349.3±11.26µs        ? ?/sec
get_or_spawn/individual                        1.13   732.1±43.35µs        ? ?/sec     1.00   645.6±41.44µs        ? ?/sec
insert_commands/insert                         1.12   623.9±37.48µs        ? ?/sec     1.00   557.4±34.99µs        ? ?/sec
insert_commands/insert_batch                   1.16   401.4±17.00µs        ? ?/sec     1.00   347.4±12.87µs        ? ?/sec
insert_simple/base                             1.08    416.9±5.60µs        ? ?/sec     1.00    385.2±4.14µs        ? ?/sec
insert_simple/unbatched                        1.06   934.5±44.58µs        ? ?/sec     1.00   881.3±47.86µs        ? ?/sec
spawn_commands/2000_entities                   1.09   190.7±11.41µs        ? ?/sec     1.00    174.7±9.15µs        ? ?/sec
spawn_commands/4000_entities                   1.10   386.5±25.33µs        ? ?/sec     1.00   352.3±18.81µs        ? ?/sec
spawn_commands/6000_entities                   1.10   586.2±34.42µs        ? ?/sec     1.00   535.3±27.25µs        ? ?/sec
spawn_commands/8000_entities                   1.08   778.5±45.15µs        ? ?/sec     1.00   718.0±33.66µs        ? ?/sec
spawn_world/10000_entities                     1.04  1026.4±195.46µs        ? ?/sec    1.00  985.8±253.37µs        ? ?/sec
spawn_world/1000_entities                      1.06   103.8±20.23µs        ? ?/sec     1.00    97.6±18.22µs        ? ?/sec
spawn_world/100_entities                       1.15     11.4±4.25µs        ? ?/sec     1.00      9.9±1.87µs        ? ?/sec
spawn_world/10_entities                        1.05  1030.8±229.78ns        ? ?/sec    1.00  986.2±231.12ns        ? ?/sec
spawn_world/1_entities                         1.01   105.1±23.33ns        ? ?/sec     1.00   104.6±31.84ns        ? ?/sec
```

---

## Changelog
Changed: `Bundle::get_components` now takes a `FnMut(StorageType, OwningPtr)`. The provided storage type must be correct for the component being fetched.

Author: james7132

crates/bevy_ecs/macros/src/lib.rs

@@ -170,7 +170,10 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
     let struct_name = &ast.ident;
 
     TokenStream::from(quote! {
-        /// SAFETY: ComponentId is returned in field-definition-order. [from_components] and [get_components] use field-definition-order
+        // SAFETY:
+        // - ComponentId is returned in field-definition-order. [from_components] and [get_components] use field-definition-order
+        // - `Bundle::get_components` is exactly once for each member. Rely's on the Component -> Bundle implementation to properly pass
+        //   the correct `StorageType` into the callback.
         unsafe impl #impl_generics #ecs_path::bundle::Bundle for #struct_name #ty_generics #where_clause {
             fn component_ids(
                 components: &mut #ecs_path::component::Components,
@@ -191,7 +194,11 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             }
 
             #[allow(unused_variables)]
-            fn get_components(self, func: &mut impl FnMut(#ecs_path::ptr::OwningPtr<'_>)) {
+            #[inline]
+            fn get_components(
+                self,
+                func: &mut impl FnMut(#ecs_path::component::StorageType, #ecs_path::ptr::OwningPtr<'_>)
+            ) {
                 #(#field_get_components)*
             }
         }

crates/bevy_ecs/src/bundle.rs

@@ -9,7 +9,7 @@ use crate::{
         Archetype, ArchetypeId, ArchetypeRow, Archetypes, BundleComponentStatus, ComponentStatus,
         SpawnBundleStatus,
     },
-    component::{Component, ComponentId, Components, StorageType, Tick},
+    component::{Component, ComponentId, ComponentStorage, Components, StorageType, Tick},
     entity::{Entities, Entity, EntityLocation},
     storage::{SparseSetIndex, SparseSets, Storages, Table, TableRow},
 };
@@ -130,7 +130,6 @@ use std::{any::TypeId, collections::HashMap};
 /// That is, there is no safe way to implement this trait, and you must not do so.
 /// If you want a type to implement [`Bundle`], you must use [`derive@Bundle`](derive@Bundle).
 ///
-///
 /// [`Query`]: crate::system::Query
 // Some safety points:
 // - [`Bundle::component_ids`] must return the [`ComponentId`] for each component type in the
@@ -159,15 +158,19 @@ pub unsafe trait Bundle: Send + Sync + 'static {
         F: for<'a> FnMut(&'a mut T) -> OwningPtr<'a>,
         Self: Sized;
 
+    // SAFETY:
+    // The `StorageType` argument passed into [`Bundle::get_components`] must be correct for the
+    // component being fetched.
+    //
     /// Calls `func` on each value, in the order of this bundle's [`Component`]s. This passes
     /// ownership of the component values to `func`.
     #[doc(hidden)]
-    fn get_components(self, func: &mut impl FnMut(OwningPtr<'_>));
+    fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>));
 }
 
 // SAFETY:
 // - `Bundle::component_ids` calls `ids` for C's component id (and nothing else)
-// - `Bundle::get_components` is called exactly once for C.
+// - `Bundle::get_components` is called exactly once for C and passes the component's storage type based on it's associated constant.
 // - `Bundle::from_components` calls `func` exactly once for C, which is the exact value returned by `Bundle::component_ids`.
 unsafe impl<C: Component> Bundle for C {
     fn component_ids(
@@ -188,8 +191,9 @@ unsafe impl<C: Component> Bundle for C {
         func(ctx).read()
     }
 
-    fn get_components(self, func: &mut impl FnMut(OwningPtr<'_>)) {
-        OwningPtr::make(self, func);
+    #[inline]
+    fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) {
+        OwningPtr::make(self, |ptr| func(C::Storage::STORAGE_TYPE, ptr));
     }
 }
 
@@ -199,6 +203,8 @@ macro_rules! tuple_impl {
         // - `Bundle::component_ids` calls `ids` for each component type in the
         // bundle, in the exact order that `Bundle::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.
         unsafe impl<$($name: Bundle),*> Bundle for ($($name,)*) {
             #[allow(unused_variables)]
             fn component_ids(components: &mut Components, storages: &mut Storages, ids: &mut impl FnMut(ComponentId)){
@@ -217,7 +223,8 @@ macro_rules! tuple_impl {
             }
 
             #[allow(unused_variables, unused_mut)]
-            fn get_components(self, func: &mut impl FnMut(OwningPtr<'_>)) {
+            #[inline(always)]
+            fn get_components(self, func: &mut impl FnMut(StorageType, OwningPtr<'_>)) {
                 #[allow(non_snake_case)]
                 let ($(mut $name,)*) = self;
                 $(
@@ -254,7 +261,6 @@ impl SparseSetIndex for BundleId {
 pub struct BundleInfo {
     pub(crate) id: BundleId,
     pub(crate) component_ids: Vec<ComponentId>,
-    pub(crate) storage_types: Vec<StorageType>,
 }
 
 impl BundleInfo {
@@ -268,11 +274,6 @@ impl BundleInfo {
         &self.component_ids
     }
 
-    #[inline]
-    pub fn storage_types(&self) -> &[StorageType] {
-        &self.storage_types
-    }
-
     pub(crate) fn get_bundle_inserter<'a, 'b>(
         &'b self,
         entities: &'a mut Entities,
@@ -386,9 +387,9 @@ impl BundleInfo {
         // 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 |component_ptr| {
+        bundle.get_components(&mut |storage_type, component_ptr| {
             let component_id = *self.component_ids.get_unchecked(bundle_component);
-            match self.storage_types[bundle_component] {
+            match storage_type {
                 StorageType::Table => {
                     let column = table.get_column_mut(component_id).unwrap();
                     // SAFETY: bundle_component is a valid index for this bundle
@@ -402,8 +403,11 @@ impl BundleInfo {
                     }
                 }
                 StorageType::SparseSet => {
-                    let sparse_set = sparse_sets.get_mut(component_id).unwrap();
-                    sparse_set.insert(entity, component_ptr, change_tick);
+                    sparse_sets.get_mut(component_id).unwrap().insert(
+                        entity,
+                        component_ptr,
+                        change_tick,
+                    );
                 }
             }
             bundle_component += 1;
@@ -707,10 +711,9 @@ impl Bundles {
             let mut component_ids = Vec::new();
             T::component_ids(components, storages, &mut |id| component_ids.push(id));
             let id = BundleId(bundle_infos.len());
-            // SAFETY: T::component_id ensures info was created
-            let bundle_info = unsafe {
-                initialize_bundle(std::any::type_name::<T>(), component_ids, id, components)
-            };
+            let bundle_info =
+                // SAFETY: T::component_id ensures info was created
+                unsafe { initialize_bundle(std::any::type_name::<T>(), component_ids, id) };
             bundle_infos.push(bundle_info);
             id
         });
@@ -726,16 +729,7 @@ unsafe fn initialize_bundle(
     bundle_type_name: &'static str,
     component_ids: Vec<ComponentId>,
     id: BundleId,
-    components: &mut Components,
 ) -> BundleInfo {
-    let mut storage_types = Vec::new();
-
-    for &component_id in &component_ids {
-        // SAFETY: component_id exists and is therefore valid
-        let component_info = components.get_info_unchecked(component_id);
-        storage_types.push(component_info.storage_type());
-    }
-
     let mut deduped = component_ids.clone();
     deduped.sort();
     deduped.dedup();
@@ -745,9 +739,5 @@ unsafe fn initialize_bundle(
         bundle_type_name
     );
 
-    BundleInfo {
-        id,
-        component_ids,
-        storage_types,
-    }
+    BundleInfo { id, component_ids }
 }