ecs: fix World::insert, finish a TODO related to it

The TODO was that if the archetype has a single entity, add a component column for the new component instead of moving the entity to a brand new archetype

lyra-ecs/src/archetype.rs

@@ -55,7 +55,6 @@ impl ComponentColumn {
     /// Set a component from pointer at an entity index.
     /// 
     /// # Safety
-    /// 
     /// This column must have space to fit the component, if it does not have room it will panic.
     pub unsafe fn set_at(&mut self, entity_index: usize, comp_src: NonNull<u8>, tick: Tick) {
         assert!(entity_index < self.capacity);
@@ -76,6 +75,23 @@ impl ComponentColumn {
         }
     }
 
+    /// Inserts an entity and its component at a specific index.
+    pub unsafe fn insert_entity(&mut self, entity_index: usize, comp_src: NonNull<u8>, tick: Tick) {
+        assert!(entity_index < self.capacity);
+
+        let mut data = self.data.borrow_mut();
+        let data = data.deref_mut();
+        
+        let size = self.info.layout().size();
+        let dest = NonNull::new_unchecked(data.as_ptr().add(entity_index * size));
+        ptr::copy_nonoverlapping(comp_src.as_ptr(), dest.as_ptr(), size);
+
+        // check if a component spot is being set twice and that the entity's tick is
+        // already stored
+        self.entity_ticks.push(tick);
+        self.len += 1;
+    }
+
     /// Get a component at an entities index.
     /// 
     /// # Safety
@@ -153,6 +169,8 @@ impl ComponentColumn {
     pub unsafe fn remove_component(&mut self, entity_index: usize, tick: &Tick) -> Option<usize> {
         let _ = tick; // may be used at some point
 
+        debug_assert!(self.len > 0, "There are no entities in the Archetype to remove from!");
+        
         let mut data = self.data.borrow_mut();
         let data = data.deref_mut();
 
@@ -170,8 +188,7 @@ impl ComponentColumn {
             mem::swap(&mut old_comp_ptr, &mut new_comp_ptr); // new_comp_ptr is now the old ptr
 
             // make sure to keep entity indexes correct in the ticks list as well
-            self.entity_ticks.swap(moved_index, entity_index);
-            self.entity_ticks.pop();
+            self.entity_ticks.swap_remove(entity_index);
 
             Some(moved_index)
         } else { None };
@@ -283,8 +300,7 @@ impl Archetype {
 
         bundle.take(|data, type_id, _size| {
             let col = self.get_column_mut(type_id).unwrap();
-            unsafe { col.set_at(entity_index.0 as usize, data, *tick); }
-            col.len += 1;
+            unsafe { col.insert_entity(entity_index.0 as usize, data, *tick); }
         });
         
         entity_index
@@ -293,21 +309,21 @@ impl Archetype {
     /// Removes an entity from the Archetype and frees its components. Returns the entity record
     /// that took its place in the component column.
     pub(crate) fn remove_entity(&mut self, entity: Entity, tick: &Tick) -> Option<(Entity, ArchetypeEntityId)> {
-        let entity_index = *self.entity_ids.get(&entity)
+        let entity_index = self.entity_ids.remove(&entity)
             .expect("The entity is not in this Archetype!");
         let mut removed_entity: Option<(Entity, ArchetypeEntityId)> = None;
 
         for c in self.columns.iter_mut() {
             let moved_entity = unsafe { c.remove_component(entity_index.0 as usize, tick) };
 
-            if let Some(res) = moved_entity {
+            if let Some(moved_idx) = moved_entity {
                 if let Some((_, aid)) = removed_entity {
                     // Make sure that the moved entity is the same as what was moved in other columns.
-                    assert!(res as u64 == aid.0);
+                    assert!(moved_idx as u64 == aid.0);
                 } else {
-                    // This is the first move, so find the EntityId that points to the column index.
-                    let just_removed = self.entities[res];
-                    removed_entity = Some((just_removed, ArchetypeEntityId(res as u64)));
+                    // This is the first move, so find the Entity that was moved into this index.
+                    let just_removed = self.entities[moved_idx];
+                    removed_entity = Some((just_removed, ArchetypeEntityId(moved_idx as u64)));
                 }
             } else {
                 // If there wasn't a moved entity, make sure no other columns moved something. 
@@ -316,12 +332,15 @@ impl Archetype {
         }
 
         // safe from the .expect at the start of this method.
-        self.entity_ids.remove(&entity).unwrap();
-        if self.entities.len() > 1 {
-            let len = self.entities.len();
-            self.entities.swap(entity_index.0 as _,  len - 1);
+        //self.entity_ids.remove(&entity).unwrap();
+        
+        // update the archetype index of the moved entity
+        if let Some((moved, _old_idx)) = removed_entity {
+            self.entity_ids.insert(moved, entity_index);
         }
-        self.entities.pop().unwrap();
+        
+        let removed = self.entities.swap_remove(entity_index.0 as _);
+        assert_eq!(removed, entity);
         
         // now change the ArchetypeEntityId to be the index that the moved entity was moved into.
         removed_entity.map(|(e, _a)| (e, entity_index))
@@ -390,7 +409,8 @@ impl Archetype {
             self.capacity = new_cap;
         }
 
-        debug_assert_eq!(self.entity_ids.len(), self.entities.len(), "Somehow the Archetype's entity storage got unsynced");
+        debug_assert_eq!(self.entity_ids.len(), self.entities.len(),
+            "Somehow the Archetype's entity storage got unsynced");
         let entity_index = ArchetypeEntityId(self.entity_ids.len() as u64);
         self.entity_ids.insert(entity, entity_index);
         self.entities.push(entity);
@@ -402,6 +422,12 @@ impl Archetype {
         entity_index
     }
 
+    /// Ensure that the internal entity lists are synced in length
+    pub(crate) fn ensure_synced(&self) {
+        debug_assert_eq!(self.entity_ids.len(), self.entities.len(),
+            "Somehow the Archetype's entity storage got unsynced");
+    }
+
     /// Moves the entity from this archetype into another one.
     /// 
     /// # Safety
@@ -453,6 +479,36 @@ impl Archetype {
     pub fn has_entity(&self, e: Entity) -> bool {
         self.entity_ids.contains_key(&e)
     }
+
+    /// Extend the Archetype by adding more columns.
+    /// 
+    /// In order to extend the Archetype, the archetype needs the components for the entities
+    /// it already has. These are provided through the `new_columns` parameter. **If the Vec
+    /// does not have the same amount of bundles in it as the amount of entities in the
+    /// Archetype, it will panic!**
+    pub fn extend<B: Bundle>(&mut self, tick: &Tick, new_columns: Vec<B>) {
+        debug_assert_eq!(new_columns.len(), self.len(), "The amount of provided column does not \
+            match the amount of entities");
+        
+        let column_info = new_columns.iter()
+            .next()
+            .unwrap()
+            .info();
+
+        for coli in column_info.into_iter() {
+            let col = unsafe { ComponentColumn::new(coli, self.capacity) };
+            self.columns.push(col);
+        }
+
+        for (eid, bundle) in new_columns.into_iter().enumerate() {
+            bundle.take(|ptr, tyid, _size| {
+                unsafe {
+                    let col = self.get_column_mut(tyid).unwrap();
+                    col.insert_entity(eid, ptr, tick.clone());
+                }
+            });
+        }
+    }
 }
 
 #[cfg(test)]
@@ -666,4 +722,24 @@ mod tests {
         assert_eq!(unsafe { *ptr.cast::<u32>().as_ref() }, comp);
         assert_eq!(col.info, info);
     }
+
+    /// Tests removing an entity from the Archetype when it is the only entity in it.
+    #[test]
+    fn remove_single_entity() {
+        let info = (Vec2::new(0.0, 0.0),).info();
+        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), info);
+
+        let ae = Entity {
+            id: EntityId(0),
+            generation: 0
+        };
+
+        a.add_entity(
+            ae,
+            Vec2::new(10.0, 50.0),
+            &Tick::default()
+        );
+
+        a.remove_entity(ae, &Tick::default());
+    }
 }

lyra-ecs/src/lib.rs

@@ -30,7 +30,7 @@ pub use component::*;
 pub mod query;
 //pub use query::*;
 
-mod relation;
+pub mod relation;
 pub use relation::Relation;
 
 mod component_info;

lyra-ecs/src/relation/child_of.rs

@@ -0,0 +1,8 @@
+use super::Relation;
+
+// TODO: Delete child entities when the parent is deleted
+pub struct ChildOf;
+
+impl Relation for ChildOf {
+    
+}

lyra-ecs/src/relation/mod.rs

@@ -10,19 +10,20 @@ use crate::lyra_engine;
 use crate::World;
 
 mod relates_to;
-#[doc(hidden)]
 pub use relates_to::*;
 
 mod relate_pair;
-#[doc(hidden)]
 pub use relate_pair::*;
 
+mod child_of;
+pub use child_of::*;
+
 #[allow(unused_variables)]
 pub trait Relation: 'static {
     /// called when a relation of this type is set on a target
-    fn relation_add(&self, origin: Entity, target: Entity) { }
+    fn relation_add(&mut self, world: &mut World, origin: Entity, target: Entity) { }
     /// called when a relation is removed
-    fn relation_remove(&self, origin: Entity, target: Entity) { }
+    fn relation_remove(&mut self, world: &mut World, origin: Entity, target: Entity) { }
 }
 
 /// A component that stores the target of a relation.
@@ -31,10 +32,16 @@ pub trait Relation: 'static {
 /// entities that the relation targets.
 #[derive(Component)]
 pub struct RelationOriginComponent<R: Relation> {
-    pub(crate) relation: R,
+    pub relation: R,
     target: Entity,
 }
 
+impl<R: Relation> RelationOriginComponent<R> {
+    pub fn target(&self) -> Entity {
+        self.target
+    }
+}
+
 /// A component that stores the origin of a relation.
 /// 
 /// This component is on the target of the relation and can be used to find the
@@ -79,12 +86,12 @@ impl World {
         };
         self.insert(target, comp);
 
-        let comp = RelationOriginComponent {
+        let mut comp = RelationOriginComponent {
             relation,
             target,
         };
         
-        comp.relation.relation_add(origin, target);
+        comp.relation.relation_add(self, origin, target);
         self.insert(origin, comp);
     }
 }

lyra-ecs/src/relation/relate_pair.rs

@@ -69,7 +69,7 @@ where
 
     unsafe fn fetch<'a>(&self, _world: &'a World, archetype: &'a crate::archetype::Archetype, tick: crate::Tick) -> Self::Fetch<'a> {
         let _ = tick;
-        let col = archetype.get_column(self.type_id())
+        let col = archetype.get_column(TypeId::of::<RelationOriginComponent<R>>())
             .expect("You ignored 'can_visit_archetype'!");
 
         FetchRelatePair {

lyra-ecs/src/world.rs

@@ -133,8 +133,6 @@ impl World {
     where
         B: Bundle
     {
-        // TODO: If the archetype has a single entity, add a component column for the new
-        // component instead of moving the entity to a brand new archetype.
         // TODO: If the entity already has the components in `bundle`, update the values of the
         // components with the bundle.
 
@@ -142,6 +140,7 @@ impl World {
 
         let record = self.entities.entity_record(entity).unwrap();
         let current_arch = self.archetypes.get(&record.id).unwrap();
+        let current_arch_len = current_arch.len();
 
         let mut col_types: Vec<DynTypeId> = current_arch.columns.iter().map(|c| c.info.type_id()).collect();
         let orig_col = col_types.clone();
@@ -150,7 +149,9 @@ impl World {
         let mut col_infos: Vec<ComponentInfo> = current_arch.columns.iter().map(|c| c.info).collect();
         col_infos.extend(bundle.info());
 
-        let col_ptrs: Vec<(NonNull<u8>, ComponentInfo)> = current_arch.columns.iter().map(|c| unsafe { (NonNull::new_unchecked(c.borrow_ptr().as_ptr()), c.info) }).collect();
+        let col_ptrs: Vec<(NonNull<u8>, ComponentInfo)> = current_arch.columns.iter()
+            .map(|c| unsafe { (NonNull::new_unchecked(c.borrow_ptr().as_ptr()), c.info) })
+            .collect();
 
         if let Some(arch) = self.archetypes.values_mut().find(|a| a.is_archetype_for(&col_types)) {
             let res_index = arch.reserve_one(entity);
@@ -158,19 +159,21 @@ impl World {
             for (col_type, (col_ptr, col_info)) in orig_col.into_iter().zip(col_ptrs.into_iter()) {
                 unsafe {
                     let ptr = NonNull::new_unchecked(col_ptr.as_ptr()
-                        .add(res_index.0 as usize * col_info.layout().size()));
+                        .add(record.index.0 as usize * col_info.layout().size()));
                     let col = arch.get_column_mut(col_type).unwrap();
+                    // set_at is used since the entity was reserved
                     col.set_at(res_index.0 as _, ptr, tick);
                 }
             }
 
             bundle.take(|data, type_id, _size| {
                 let col = arch.get_column_mut(type_id).unwrap();
+                // set_at is used since the entity was reserved
                 unsafe { col.set_at(res_index.0 as _, data, tick); }
-                col.len += 1;
             });
 
             arch.entity_ids.insert(entity, res_index);
+            arch.ensure_synced();
             
             let new_record = Record {
                 id: arch.id(),
@@ -178,10 +181,30 @@ impl World {
             };
             self.entities.insert_entity_record(entity, new_record);
         } else {
+            if current_arch_len == 1 {
+                // if this entity is the only entity for this archetype, add more columns to it
+                let current_arch = self.archetypes.get_mut(&record.id).unwrap();
+                current_arch.extend(&tick, vec![bundle]);
+
+                return;
+            }
+
             let new_arch_id = self.next_archetype_id.increment();
             let mut archetype = Archetype::from_bundle_info(new_arch_id, col_infos);
             let entity_arch_id = archetype.add_entity(entity, bundle, &tick);
             
+            // move the old components into the new archetype
+            for (column_ptr, column_info) in col_ptrs.into_iter() {
+                unsafe {
+                    // ptr of component for the entity
+                    let comp_ptr = NonNull::new_unchecked(column_ptr.as_ptr()
+                        .add(record.index.0 as usize * column_info.layout().size()));
+                    
+                    let col = archetype.get_column_mut(column_info.type_id()).unwrap();
+                    col.insert_entity(entity_arch_id.0 as _, comp_ptr, tick);
+                }
+            }
+
             self.archetypes.insert(new_arch_id, archetype);
 
             // Create entity record and store it
@@ -194,7 +217,13 @@ impl World {
         }
 
         let current_arch = self.archetypes.get_mut(&record.id).unwrap();
+        if current_arch.len() > 1 {
             current_arch.remove_entity(entity, &tick);
+        } else if current_arch.len() == 1 {
+            // The old archetype will only be removed if there was another archetype that would
+            // work for the entity's components, and the old archetype only had a single entity.
+            self.archetypes.remove(&record.id).unwrap();
+        }
     }
 
     pub fn entity_archetype(&self, entity: Entity) -> Option<&Archetype> {
@@ -330,7 +359,7 @@ impl World {
 
 #[cfg(test)]
 mod tests {
-    use crate::{tests::{Vec2, Vec3}, query::TickOf};
+    use crate::{query::TickOf, tests::{Vec2, Vec3}, Entity};
 
     use super::World;
 
@@ -453,6 +482,46 @@ mod tests {
         assert!(world.view_one::<&Vec3>(e).get().is_some())
     }
 
+    #[test]
+    fn insert_multiple_times() {
+        let v2s = &[Vec2::rand(), Vec2::rand(), Vec2::rand()];
+        let v3s = &[Vec3::rand(), Vec3::rand(), Vec3::rand()];
+
+        let mut world = World::new();
+        let e1 = world.spawn(v2s[0]);
+        let e2 = world.spawn(v2s[1]);
+        let e3 = world.spawn(v2s[2]);
+        println!("Spawned entities");
+
+        let ev2 = world.view_one::<&Vec2>(e2).get()
+            .expect("Failed to find Vec2 and Vec3 on inserted entity!");
+        assert_eq!(*ev2, v2s[1]);
+        drop(ev2);
+
+        let insert_and_assert = |world: &mut World, e: Entity, v2: Vec2, v3: Vec3| {
+            println!("inserting entity");
+            world.insert(e, (v3,));
+            println!("inserted entity");
+
+            let (ev2, ev3) = world.view_one::<(&Vec2, &Vec3)>(e).get()
+                .expect("Failed to find Vec2 and Vec3 on inserted entity!");
+            assert_eq!(*ev2, v2);
+            assert_eq!(*ev3, v3);
+        };
+
+        insert_and_assert(&mut world, e2, v2s[1], v3s[1]);
+        println!("Entity 2 is good");
+        insert_and_assert(&mut world, e3, v2s[2], v3s[2]);
+        println!("Entity 3 is good");
+        assert_eq!(world.archetypes.len(), 2);
+        println!("No extra archetypes were created");
+
+        insert_and_assert(&mut world, e1, v2s[0], v3s[0]);
+        println!("Entity 1 is good");
+        assert_eq!(world.archetypes.len(), 1);
+        println!("Empty archetype was removed");
+    }
+
     #[test]
     fn view_one() {
         let v = Vec2::rand();
@@ -468,11 +537,6 @@ mod tests {
     fn view_change_tracking() {
         let mut world = World::new();
 
-        /* let v = Vec2::rand();
-        world.spawn((v,));
-        let v = Vec2::rand();
-        world.spawn((v,)); */
-
         println!("spawning");
         world.spawn((Vec2::new(10.0, 10.0),));
         world.spawn((Vec2::new(5.0, 5.0),));