lyra-engine/lyra-ecs/src/archetype.rs

use std::{ptr::{NonNull, self}, alloc::{self, Layout, alloc, dealloc}, mem, collections::HashMap, cell::{RefCell, Ref, RefMut, BorrowError, BorrowMutError}, ops::{DerefMut, Deref}};

use crate::{world::{Entity, ArchetypeEntityId}, bundle::Bundle, component_info::ComponentInfo, DynTypeId};

pub struct ComponentColumn {
    pub data: RefCell<NonNull<u8>>,
    pub len: usize,
    pub capacity: usize,
    pub info: ComponentInfo,
}

impl Drop for ComponentColumn {
    fn drop(&mut self) {
        let data = self.data.borrow_mut();
        let data = data.as_ptr();

        unsafe {
            // layout of current alloc
            let layout = Layout::from_size_align_unchecked(self.info.layout.size * self.capacity,
                self.info.layout.alignment);
            dealloc(data, layout);
        }
    }
}

#[allow(dead_code)]
impl ComponentColumn {
    /// Creates an invalid component column. Do not attempt to use it.
    /* pub unsafe fn dangling() -> Self {
        ComponentColumn {
            data: RefCell::
            capacity: 0,
            info: ComponentInfo::new::<()>(),
            entry_size: 0,
            len: 0,
        }
    } */
    
    pub unsafe fn alloc(component_layout: Layout, capacity: usize) -> NonNull<u8> {
        let new_layout = Layout::from_size_align(
            component_layout.size().checked_mul(capacity).unwrap(),
            component_layout.align()
        ).unwrap();

        if let Some(data) = NonNull::new(alloc(new_layout)) {
            data
        } else {
            alloc::handle_alloc_error(new_layout)
        }
    }

    pub unsafe fn new(info: ComponentInfo, capacity: usize) -> Self {
        let data = ComponentColumn::alloc(info.layout.into_layout_unchecked(), capacity);
        
        Self {
            data: RefCell::new(data),
            capacity,
            info,
            len: 0,
        }
    }

    /// 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>) {
        assert!(entity_index < self.capacity);

        let mut data = self.data.borrow_mut();
        let data = data.deref_mut();
        
        let dest = NonNull::new_unchecked(data.as_ptr().add(entity_index * self.info.layout.size));
        ptr::copy_nonoverlapping(comp_src.as_ptr(), dest.as_ptr(), self.info.layout.size);
        
        unsafe {
            let layout = self.info.layout.into_layout_unchecked();
            std::alloc::dealloc(comp_src.as_ptr(), layout);
        }
    }

    /// Get a component at an entities index.
    /// 
    /// # Safety
    /// 
    /// This column MUST have the entity. If it does not, it WILL NOT panic and will cause UB.
    pub unsafe fn get<T>(&self, entity_index: usize) -> &T {
        let data = self.data.borrow();
        let data = data.deref();

        let ptr = NonNull::new_unchecked(data.as_ptr()
            .add(entity_index * self.info.layout.size))
            .cast();
        &*ptr.as_ptr()
    }

    /// Get a component at an entities index.
    /// 
    /// # Safety
    /// 
    /// This column must have the entity.
    pub unsafe fn get_mut<T>(&mut self, entity_index: usize) -> &mut T {
        let mut data = self.data.borrow_mut();
        let data = data.deref_mut();

        let p = data.as_ptr()
            .cast::<T>()
            .add(entity_index * self.info.layout.size);
        &mut *p
    }

    /// Grow the column to fit `new_capacity` amount of components.
    /// 
    /// Parameters:
    /// * `new_capacity` - The new capacity of components that can fit in this column.
    /// 
    /// # Safety
    /// 
    /// Will panic if `new_capacity` is less than the current capacity of the column.
    pub unsafe fn grow(&mut self, new_capacity: usize) {
        assert!(new_capacity > self.capacity);

        let mut data = self.data.borrow_mut();
        //let data = data.deref_mut();

        let mut new_ptr = Self::alloc(self.info.layout.into_layout_unchecked(), new_capacity);

        if self.len > 0 {
            ptr::copy_nonoverlapping(data.as_ptr(), new_ptr.as_ptr(), self.len * self.info.layout.size);
        }

        // dont attempt to free if we weren't able to store anything anyway
        if self.capacity != 0 {
            let old_layout = Layout::from_size_align_unchecked(
                self.info.layout.size.checked_mul(self.capacity).unwrap(),
                self.info.layout.alignment
            );

            mem::swap(data.deref_mut(), &mut new_ptr);
            dealloc(new_ptr.as_ptr(), old_layout);
        } else {
            *data = new_ptr;
        }
        
        self.capacity = new_capacity;
    }

    /// Removes a component from the column, freeing it, and returning the old index of the entity that took its place in the column.
    pub unsafe fn remove_component(&mut self, entity_index: usize) -> Option<usize> {
        let mut data = self.data.borrow_mut();
        let data = data.deref_mut();

        let mut old_comp_ptr = NonNull::new_unchecked(data.as_ptr()
            .add(entity_index * self.info.layout.size));

        let moved_index = if entity_index != self.len - 1 {
            let moved_index = self.len - 1;
            let mut new_comp_ptr = NonNull::new_unchecked(data.as_ptr()
                .add(moved_index * self.info.layout.size));

            ptr::copy_nonoverlapping(new_comp_ptr.as_ptr(), old_comp_ptr.as_ptr(), self.info.layout.size);

            mem::swap(&mut old_comp_ptr, &mut new_comp_ptr); // new_comp_ptr is now the old ptr
            Some(moved_index)
        } else { None };

        self.len -= 1;

        moved_index
    }

    pub fn borrow_ptr<'a>(&'a self) -> Ref<'a, NonNull<u8>> {
        self.data.borrow()
    }

    pub fn borrow_mut_ptr<'a>(&'a self) -> RefMut<'a, NonNull<u8>> {
        self.data.borrow_mut()
    }

    pub fn try_borrow_ptr<'a>(&'a self) -> Result<Ref<'a, NonNull<u8>>, BorrowError> {
        self.data.try_borrow()
    }

    pub fn try_borrow_mut_ptr<'a>(&'a self) -> Result<RefMut<'a, NonNull<u8>>, BorrowMutError> {
        self.data.try_borrow_mut()
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct ArchetypeId(pub u64);

impl ArchetypeId {
    /// Increments the id and returns a new id with the value it was before incrementing.
    pub(crate) fn increment(&mut self) -> Self {
        let v = self.0;
        self.0 += 1;
        
        ArchetypeId(v)
    }
}

pub struct Archetype {
    pub id: ArchetypeId,
    pub(crate) entities: HashMap<Entity, ArchetypeEntityId>,
    pub(crate) columns: Vec<ComponentColumn>,
    pub capacity: usize,
}

/// The default capacity of the columns
const DEFAULT_CAPACITY: usize = 32;

impl Archetype {
    pub fn from_bundle_info(new_id: ArchetypeId, bundle_info: Vec<ComponentInfo>) -> Archetype {
        let columns = bundle_info.into_iter().map(|i| {
            unsafe { ComponentColumn::new(i, DEFAULT_CAPACITY) }
        }).collect();

        Archetype {
            id: new_id,
            entities: HashMap::new(),
            columns,
            capacity: DEFAULT_CAPACITY,
        }
    }

    /// Add an entity and its component bundle to the Archetype
    /// 
    /// # Safety:
    /// 
    /// Archetype must contain all of the components
    pub(crate) fn add_entity<B>(&mut self, entity: Entity, bundle: B) -> ArchetypeEntityId
    where
        B: Bundle
    {
        if self.capacity == self.entities.len() {
            let new_cap = self.capacity * 2;
            self.grow_columns(new_cap);
            self.capacity = new_cap;
        }

        let entity_index = self.entities.len();
        self.entities.insert(entity, ArchetypeEntityId(entity_index as u64));

        bundle.take(|data, type_id, _size| {
            let col = self.get_column_mut(type_id).unwrap();
            unsafe { col.set_at(entity_index, data); }
            col.len += 1;
        });
        
        ArchetypeEntityId(entity_index as u64)
    }

    /// 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) -> Option<(Entity, ArchetypeEntityId)> {
        let entity_index = *self.entities.get(&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) };

            // Make sure that the moved entity is the same as what was moved in other columns.
            // If this is the first move, find the EntityId that points to the column index.
            // If there wasn't a moved entity, make sure no other columns moved something.
            if let Some(res) = moved_entity {
                if let Some((_, aid)) = removed_entity {
                    assert!(res as u64 == aid.0); // make sure we removed the same entity
                } else {
                    let replaced_entity = self.entities.iter().find(|(_, a)| a.0 == res as u64)
                        .map(|(e, _a)| *e).expect("Failure to find entity for moved component!");
                    removed_entity = Some((replaced_entity, ArchetypeEntityId(res as u64)));
                }
            } else {
                assert!(removed_entity.is_none());
            }
        }

        // safe from the .expect at the start of this method.
        self.entities.remove(&entity).unwrap();
        
        // now change the ArchetypeEntityId to be the index that the moved entity was moved into.
        removed_entity.map(|(e, _a)| (e, entity_index))

    }

    /// Returns a boolean indicating whether this archetype can store the TypeIds given
    pub fn is_archetype_for(&self, types: &Vec<DynTypeId>) -> bool {
        if types.len() == self.columns.len() {
            self.columns.iter().all(|c| types.contains(&c.info.type_id))
        } else { false }
    }

    /// Returns a boolean indicating whether this archetype has a column for `comp_type`
    pub fn has_column(&self, comp_type: DynTypeId) -> bool {
        self.columns.iter().any(|c| comp_type == c.info.type_id)
    }

    /// Returns a boolean indicating whether this archetype is empty or not.
    pub fn is_empty(&self) -> bool {
        self.entities.is_empty()
    }

    /// Returns the amount of entities that are stored in the archetype.
    pub fn len(&self) -> usize {
        self.entities.len()
    }

    /// Grows columns in the archetype
    /// 
    /// Parameters:
    /// * `new_capacity` - The new capacity of components that can fit in this column.
    /// 
    /// # Safety
    /// 
    /// Will panic if new_capacity is less than the current capacity
    fn grow_columns(&mut self, new_capacity: usize) {
        assert!(new_capacity > self.capacity);

        for c in self.columns.iter_mut() {
            unsafe { c.grow(new_capacity); }
        }

        self.capacity = new_capacity;
    }

    pub fn get_column(&self, type_id: DynTypeId) -> Option<&ComponentColumn> {
        self.columns.iter().find(|c| c.info.type_id == type_id)
    }

    pub fn get_column_mut(&mut self, type_id: DynTypeId) -> Option<&mut ComponentColumn> {
        self.columns.iter_mut().find(|c| c.info.type_id == type_id)
    }

    /// Reserves a slot in the columns for an entity and returns the index of that reserved spot
    pub fn reserve_one(&mut self, entity: Entity) -> ArchetypeEntityId {
        if self.capacity == self.entities.len() {
            let new_cap = self.capacity * 2;
            self.grow_columns(new_cap);
            self.capacity = new_cap;
        }

        let entity_index = self.entities.len();
        self.entities.insert(entity, ArchetypeEntityId(entity_index as u64));

        for col in self.columns.iter_mut() {
            col.len += 1;
        }

        ArchetypeEntityId(entity_index as u64)
    }

    /// Moves the entity from this archetype into another one.
    /// 
    /// # Safety
    /// The entity IS NOT removed from the old archetype. You must manually call [`Archetype::remove_entity`].
    /// It was done this way because I had some borrow check issues when writing [`World::insert`]
    /// related to borrowing mutably from self more than once.
    pub fn move_into<B>(&self, into_arch: &mut Archetype, entity: Entity, new_components: B)
    where
        B: Bundle
    {
        let new_index = into_arch.reserve_one(entity);

        //let infos: Vec<ComponentInfo> = into_arch.columns.iter().map(|c| c.info).collect();

        // move the existing components into the new archetype
        for col in self.columns.iter() {
            let into_col = into_arch.get_column_mut(col.info.type_id).unwrap();
            
            // copy from the old column into the new column, then remove it from the old one
            unsafe {
                let ptr = col.borrow_ptr();
                let ptr = NonNull::new_unchecked(ptr.as_ptr()
                    .add(new_index.0 as usize * col.info.layout.size));
                into_col.set_at(new_index.0 as _, ptr);
            }
        }

        // now move the new components into the new archetype
        new_components.take(|data, type_id, _size| {
            let col = into_arch.get_column_mut(type_id).unwrap();
            unsafe { col.set_at(new_index.0 as _, data); }
            col.len += 1;
        });

        //self.remove_entity(entity);
    }

    pub fn entities(&self) -> &HashMap<Entity, ArchetypeEntityId> {
        &self.entities
    }
}

#[cfg(test)]
mod tests {
    use std::{alloc::Layout, ptr::NonNull};

    use rand::Rng;

    use crate::{ArchetypeId, tests::{Vec2, Vec3}, world::{Entity, EntityId}, bundle::Bundle, ComponentInfo, MemoryLayout, DynTypeId, DynamicBundle};

    use super::Archetype;

    #[test]
    fn one_entity_one_component() {
        let bundle = (Vec2::new(10.0, 20.0),);
        let entity = Entity {
            id: EntityId(0),
            generation: 0
        };

        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), bundle.info());
        let entity_arch_id = a.add_entity(entity, bundle);

        let col = a.columns.get(0).unwrap();
        let vec2: &Vec2 = unsafe { col.get(entity_arch_id.0 as usize) };
        assert_eq!(vec2.clone(), bundle.0);
    }

    #[test]
    fn one_entity_two_component() {
        let bundle = (Vec2::new(10.0, 20.0),Vec3::new(15.0, 54.0, 84.0));
        let entity = Entity {
            id: EntityId(0),
            generation: 0
        };

        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), bundle.info());
        let entity_arch_id = a.add_entity(entity, bundle);

        let col = a.columns.get(0).unwrap();
        let vec2: &Vec2 = unsafe { col.get(entity_arch_id.0 as usize) };
        assert_eq!(vec2.clone(), bundle.0);

        let col = a.columns.get(1).unwrap();
        let vec3: &Vec3 = unsafe { col.get(entity_arch_id.0 as usize) };
        assert_eq!(vec3.clone(), bundle.1);
    }

    #[test]
    fn two_entity_one_component() {
        let b1 = (Vec2::new(10.0, 20.0),);
        let e1 = Entity {
            id: EntityId(0),
            generation: 0
        };
        let b2 = (Vec2::new(19.0, 43.0),);
        let e2 = Entity {
            id: EntityId(1),
            generation: 0
        };

        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), b1.info());
        let earch1 = a.add_entity(e1, b1);
        let earch2 = a.add_entity(e2, b2);

        let col = a.columns.get(0).unwrap();
        let vec2: &Vec2 = unsafe { col.get(earch1.0 as usize) };
        assert_eq!(vec2.clone(), b1.0);
        let vec2: &Vec2 = unsafe { col.get(earch2.0 as usize) };
        assert_eq!(vec2.clone(), b2.0);
    }

    #[test]
    fn two_entity_two_component() {
        let b1 = (Vec2::new(10.0, 20.0), Vec3::new(84.0, 283.0, 28.0));
        let e1 = Entity {
            id: EntityId(0),
            generation: 0
        };
        let b2 = (Vec2::new(19.0, 43.0), Vec3::new(74.0, 28.0, 93.0));
        let e2 = Entity {
            id: EntityId(1),
            generation: 0
        };

        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), b1.info());
        let earch1 = a.add_entity(e1, b1);
        let earch2 = a.add_entity(e2, b2);

        let col = a.columns.get(0).unwrap();
        let vec2: &Vec2 = unsafe { col.get(earch1.0 as usize) };
        assert_eq!(vec2.clone(), b1.0);
        let vec2: &Vec2 = unsafe { col.get(earch2.0 as usize) };
        assert_eq!(vec2.clone(), b2.0);

        let col = a.columns.get(1).unwrap();
        let vec3: &Vec3 = unsafe { col.get(earch1.0 as usize) };
        assert_eq!(vec3.clone(), b1.1);
        let vec3: &Vec3 = unsafe { col.get(earch2.0 as usize) };
        assert_eq!(vec3.clone(), b2.1);
    }

    /// Tests manual growth of archetypes
    #[test]
    fn archetype_growth() {
        let info = (Vec2::new(0.0, 0.0),).info();
        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), info);
        a.grow_columns(64);
        a.grow_columns(128);
        a.grow_columns(256);
    }

    /// Tests the auto growth of archetypes when adding entities
    #[test]
    fn auto_archetype_growth() {
        let mut rng = rand::thread_rng();
        let bundle_count = rng.gen_range(50..150);
        
        let mut bundles: Vec<(Vec2,)> = vec![];
        bundles.reserve(bundle_count);

        let info = (Vec2::new(0.0, 0.0),).info();
        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), info);

        for i in 0..bundle_count {
            let c = (Vec2::rand(),);
            bundles.push(c);
            
            a.add_entity(
                Entity {
                    id: EntityId(i as u64),
                    generation: 0
                },
                c
            );
        }
        println!("Inserted {} entities", bundle_count);

        let col = a.columns.get(0).unwrap();
        for i in 0..bundle_count {
            let vec2: &Vec2 = unsafe { col.get(i) };
            assert_eq!(vec2.clone(), bundles[i].0);
        }
    }

    #[test]
    fn remove_entity() {
        let bundles = vec![Vec2::rand(), Vec2::rand(), Vec2::rand()];
        println!("Bundles: {:?}", bundles);

        let info = (Vec2::new(0.0, 0.0),).info();
        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), info);

        // add the entities to the archetype
        for i in 0..bundles.len() {
            a.add_entity(
                Entity {
                    id: EntityId(i as u64),
                    generation: 0
                },
                (bundles[i],),
            );
        }

        // Remove the 'middle' entity in the column
        let moved_entity = a.remove_entity(
            Entity {
                id: EntityId(1u64),
                generation: 0
            }
        ).expect("No entity was moved");

        // The last entity in the column should have been moved
        assert!(moved_entity.0.id.0 == 2);
        assert!(moved_entity.1.0 == 1);

        // make sure that the entities' component was actually moved in the column
        let col = &a.columns[0];
        let comp = unsafe { col.get::<Vec2>(1) };
        assert_eq!(comp.clone(), bundles[2]);
    }

    /// This test simulates an archetype that stores types that rust does not know about.
    #[test]
    fn dynamic_archetype() {
        let layout = MemoryLayout::from(Layout::new::<u32>());
        let info = ComponentInfo::new_unknown(DynTypeId::Unknown(100), "u32", layout);
        let infos = vec![info.clone()];

        let mut a = Archetype::from_bundle_info(super::ArchetypeId(0), infos);

        let mut dynamic_bundle = DynamicBundle::default();
        let comp = 50u32;
        let ptr = NonNull::from(&comp).cast::<u8>();
        dynamic_bundle.push_unknown(ptr, info.clone());

        a.add_entity(
            Entity {
                id: EntityId(0),
                generation: 0
            }, dynamic_bundle
        );

        let col = a.columns.iter().next().unwrap();
        let ptr = col.borrow_ptr();
        assert_eq!(unsafe { *ptr.cast::<u32>().as_ref() }, comp);
        assert_eq!(col.info, info);
    }

    #[test]
    fn move_archetypes() {
        let bundles = vec![Vec2::rand(), ];

        let mut info = (bundles[0],).info();
        let mut a = Archetype::from_bundle_info(ArchetypeId(0), info.clone());

        let entity = Entity {
            id: EntityId(0),
            generation: 0
        };

        a.add_entity(entity, (bundles[0],));

        info.extend((Vec3::new(0.0, 0.0, 0.0),).info().into_iter());
        let mut new_a = Archetype::from_bundle_info(ArchetypeId(1), info);

        let inserting_vec3 = Vec3::rand();
        let new_bundle = (inserting_vec3,);
        a.move_into(&mut new_a, entity, new_bundle);

        let new_index = new_a.entities.get(&entity).unwrap();

        // Check Vec3
        {
            let new_col = new_a.get_column(DynTypeId::of::<Vec3>()).unwrap();
            let from_col = unsafe { new_col.get::<Vec3>(new_index.0 as _) };
            println!("expected {:?}, got {:?}", inserting_vec3, *from_col);
            assert_eq!(*from_col, inserting_vec3);
        }

        // Check Vec2
        {
            let new_col = new_a.get_column(DynTypeId::of::<Vec2>()).unwrap();
            let from_col = unsafe { new_col.get::<Vec2>(new_index.0 as _) };
            println!("expected {:?}, got {:?}", bundles[0], *from_col);
            assert_eq!(*from_col, bundles[0]);
        }

        assert!(a.entities.contains_key(&entity));
        assert!(a.remove_entity(entity).is_none()); // No entity would take its place
    }
}