1 files changed, 31 insertions, 11 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 00ba330..041c124 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -87,8 +87,12 @@ impl Storage
         &mut self,
         entity_uid: EntityUid,
         mut components: Vec<Box<dyn Component>>,
-    ) -> (ArchetypeId, EntityUid)
+    ) -> Result<(ArchetypeId, EntityUid), Error>
     {
+        if self.entity_archetype_lookup.contains_key(&entity_uid) {
+            return Err(Error::EntityAlreadyExists(entity_uid));
+        }
+
         components.sort_by_key(|component| component.id());
 
         #[cfg(feature = "debug")]
@@ -132,7 +136,7 @@ impl Storage
         self.entity_archetype_lookup
             .insert(entity_uid, archetype_id);
 
-        (archetype_id, entity_uid)
+        Ok((archetype_id, entity_uid))
     }
 
     pub fn add_components_to_entity(
@@ -171,6 +175,8 @@ impl Storage
 
         let entity = archetype.take_entity(entity_uid)?;
 
+        self.entity_archetype_lookup.remove(&entity_uid);
+
         self.push_entity(
             entity_uid,
             entity
@@ -179,7 +185,8 @@ impl Storage
                 .map(|component| component.component.into_inner())
                 .chain(components)
                 .collect(),
-        );
+        )
+        .expect("Not supposed to return Err since the entity is removed");
 
         Some(())
     }
@@ -201,6 +208,8 @@ impl Storage
 
         let component_ids_set = component_ids.into_iter().collect::<HashSet<_>>();
 
+        self.entity_archetype_lookup.remove(&entity_uid);
+
         self.push_entity(
             entity_uid,
             entity
@@ -209,7 +218,8 @@ impl Storage
                 .map(|component| component.component.into_inner())
                 .filter(|component| !component_ids_set.contains(&component.id()))
                 .collect(),
-        );
+        )
+        .expect("Not supposed to return Err since the entity is removed");
 
         Some(())
     }
@@ -315,6 +325,14 @@ impl Storage
     }
 }
 
+/// Component storage error
+#[derive(Debug, Clone, thiserror::Error)]
+pub enum Error
+{
+    #[error("Entity already exists")]
+    EntityAlreadyExists(EntityUid),
+}
+
 impl TypeName for Storage
 {
     fn type_name(&self) -> &'static str
@@ -549,13 +567,15 @@ mod tests
     {
         let mut component_storage = Storage::default();
 
-        component_storage.push_entity(
-            EntityUid::new_unique(),
-            vec![
-                Box::new(HealthPotion { _hp_restoration: 12 }),
-                Box::new(Hookshot { _range: 50 }),
-            ],
-        );
+        component_storage
+            .push_entity(
+                EntityUid::new_unique(),
+                vec![
+                    Box::new(HealthPotion { _hp_restoration: 12 }),
+                    Box::new(Hookshot { _range: 50 }),
+                ],
+            )
+            .expect("Expected Ok");
 
         assert_eq!(component_storage.archetypes.len(), 1);