From 31f8d8760b9a53d4f517c305d31953f9ce592535 Mon Sep 17 00:00:00 2001
From: bjorn3 <bjorn3@users.noreply.github.com>
Date: Tue, 13 Apr 2021 15:29:01 +0200
Subject: [PATCH] Extract monomorphic get_insert_bundle_info function

---
 crates/bevy_ecs/src/world/entity_ref.rs | 38 +++++++++++++++++++------
 1 file changed, 29 insertions(+), 9 deletions(-)

diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
index 11a7d4179e63f..f8cb1dc60a183 100644
--- a/crates/bevy_ecs/src/world/entity_ref.rs
+++ b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -2,7 +2,7 @@ use crate::{
     archetype::{Archetype, ArchetypeId, Archetypes, ComponentStatus},
     bundle::{Bundle, BundleInfo},
     component::{Component, ComponentId, ComponentTicks, Components, StorageType},
-    entity::{Entity, EntityLocation},
+    entity::{Entities, Entity, EntityLocation},
     storage::{SparseSet, Storages},
     world::{Mut, World},
 };
@@ -182,7 +182,6 @@ impl<'w> EntityMut<'w> {
             })
     }
 
-    // TODO: factor out non-generic part to cut down on monomorphization (just check perf)
     // TODO: move relevant methods to World (add/remove bundle)
     pub fn insert_bundle<T: Bundle>(&mut self, bundle: T) -> &mut Self {
         let entity = self.entity;
@@ -195,19 +194,28 @@ impl<'w> EntityMut<'w> {
         let bundle_info = self.world.bundles.init_info::<T>(components);
         let current_location = self.location;
 
-        let (archetype, bundle_status, archetype_index) = unsafe {
+        // Use a non-generic function to cut down on monomorphization
+        unsafe fn get_insert_bundle_info<'a>(
+            entities: &mut Entities,
+            archetypes: &'a mut Archetypes,
+            components: &mut Components,
+            storages: &mut Storages,
+            bundle_info: &BundleInfo,
+            current_location: EntityLocation,
+            entity: Entity,
+        ) -> (&'a Archetype, &'a Vec<ComponentStatus>, EntityLocation) {
             // SAFE: component ids in `bundle_info` and self.location are valid
             let new_archetype_id = add_bundle_to_archetype(
                 archetypes,
                 storages,
                 components,
-                self.location.archetype_id,
+                current_location.archetype_id,
                 bundle_info,
             );
             if new_archetype_id == current_location.archetype_id {
                 let archetype = &archetypes[current_location.archetype_id];
                 let edge = archetype.edges().get_add_bundle(bundle_info.id).unwrap();
-                (archetype, &edge.bundle_status, current_location.index)
+                (archetype, &edge.bundle_status, current_location)
             } else {
                 let (old_table_row, old_table_id) = {
                     let old_archetype = &mut archetypes[current_location.archetype_id];
@@ -241,22 +249,34 @@ impl<'w> EntityMut<'w> {
                     new_location
                 };
 
-                self.location = new_location;
-                entities.meta[self.entity.id as usize].location = new_location;
+                entities.meta[entity.id as usize].location = new_location;
                 let (old_archetype, new_archetype) =
                     archetypes.get_2_mut(current_location.archetype_id, new_archetype_id);
                 let edge = old_archetype
                     .edges()
                     .get_add_bundle(bundle_info.id)
                     .unwrap();
-                (&*new_archetype, &edge.bundle_status, new_location.index)
+                (&*new_archetype, &edge.bundle_status, new_location)
 
                 // Sparse set components are intentionally ignored here. They don't need to move
             }
+        }
+
+        let (archetype, bundle_status, new_location) = unsafe {
+            get_insert_bundle_info(
+                entities,
+                archetypes,
+                components,
+                storages,
+                bundle_info,
+                current_location,
+                entity,
+            )
         };
+        self.location = new_location;
 
         let table = &storages.tables[archetype.table_id()];
-        let table_row = archetype.entity_table_row(archetype_index);
+        let table_row = archetype.entity_table_row(new_location.index);
         // SAFE: table row is valid
         unsafe {
             bundle_info.write_components(