perf(ecs): use swap removal to remove entities from archetypes

1 files changed, 20 insertions, 11 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 53b1168..cb219f1 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -127,10 +127,6 @@ impl Storage
 
         archetype.push_entity(entity_uid, components);
 
-        archetype
-            .entity_lookup
-            .insert(entity_uid, archetype.entities.len() - 1);
-
         self.entity_archetype_lookup
             .insert(entity_uid, archetype_id);
 
@@ -386,6 +382,7 @@ pub struct Archetype
 {
     component_ids: HashMap<Uid, usize>,
     entity_lookup: HashMap<Uid, usize>,
+    entity_uid_lookup: Vec<Uid>,
     entities: Vec<ArchetypeEntity>,
 }
 
@@ -400,6 +397,7 @@ impl Archetype
                 .map(|(index, component_id)| (component_id, index))
                 .collect(),
             entity_lookup: HashMap::new(),
+            entity_uid_lookup: Vec::new(),
             entities: Vec::new(),
         }
     }
@@ -458,21 +456,32 @@ impl Archetype
             uid: entity_uid,
             components: components.into_iter().map(Into::into).collect(),
         });
+
+        let index = self.entities.len() - 1;
+
+        self.entity_lookup.insert(entity_uid, index);
+
+        self.entity_uid_lookup.push(entity_uid);
     }
 
     pub fn take_entity(&mut self, entity_uid: Uid) -> Option<ArchetypeEntity>
     {
         let entity_index = self.entity_lookup.remove(&entity_uid)?;
 
-        let entity = self.entities.remove(entity_index);
+        let last_entity_uid = *self
+            .entity_uid_lookup
+            .get(self.entities.len() - 1)
+            .expect("Entity UID lookup contains too few entity UIDS");
 
-        for index in self.entity_lookup.values_mut() {
-            if *index > entity_index {
-                *index -= 1;
-            }
-        }
+        // By using swap_remove, no memory reallocation occurs and only one index in the
+        // entity lookup needs to be updated
+        let removed_entity = self.entities.swap_remove(entity_index);
+
+        self.entity_lookup.insert(last_entity_uid, entity_index);
+
+        self.entity_uid_lookup.swap_remove(entity_index);
 
-        Some(entity)
+        Some(removed_entity)
     }
 }