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::ArchetypeEntityId, bundle::Bundle, component_info::ComponentInfo, DynTypeId, Tick, Entity};

#[derive(Clone)]
pub struct ComponentColumn {
    pub data: RefCell<NonNull<u8>>,
    pub len: usize,
    pub capacity: usize,
    pub info: ComponentInfo,
    /// The last tick that this component was modified at.
    pub entity_ticks: Vec<Tick>,
}

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 {
    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().unwrap(), capacity);
        
        Self {
            data: RefCell::new(data),
            capacity,
            info,
            len: 0,
            entity_ticks: Vec::new(),
        }
    }

    /// 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);

        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);

        // check if a component spot is being set twice and that the entity's tick is
        // already stored
        if self.entity_ticks.len() > entity_index {
            self.entity_ticks[entity_index].tick_to(&tick);
        } else {
            self.entity_ticks.push(tick);
        }
    }

    /// 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, tick: &Tick) -> &mut T {
        self.entity_ticks[entity_index].tick_to(tick);

        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.
    /// 
    /// Note: This does not modify the Tick of this column, since no components were actually modified.
    /// 
    /// # 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 mut new_ptr = Self::alloc(self.info.layout.into_layout().unwrap(), 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, tick: &Tick) -> Option<usize> {
        let _ = tick; // may be used at some point
        
        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

            // 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();

            Some(moved_index)
        } else { None };

        self.len -= 1;

        moved_index
    }

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

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

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

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

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)]
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)
    }
}

#[derive(Clone, Default)]
pub struct Archetype {
    pub id: ArchetypeId,
    pub(crate) entities: HashMap<Entity, ArchetypeEntityId>,
    /// map an Archetype entity id to an entity
    pub(crate) ids_to_entity: HashMap<ArchetypeEntityId, Entity>,
    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(),
            ids_to_entity: 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, tick: &Tick) -> 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 = ArchetypeEntityId(self.entities.len() as u64);
        self.entities.insert(entity, entity_index);
        self.ids_to_entity.insert(entity_index, entity);

        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;
        });
        
        entity_index
    }

    /// 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.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, tick) };

            // 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 all columns 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();
        self.ids_to_entity.remove(&entity_index).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)
    }

    /// Returns a mutable borrow to a component column for `type_id`.
    /// 
    /// Note: This does not modify the tick for the column!
    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 = ArchetypeEntityId(self.entities.len() as u64);
        self.entities.insert(entity, entity_index);
        self.ids_to_entity.insert(entity_index, entity);

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

        entity_index
    }

    /// 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
    }

    pub fn entity_of_index(&self, id: ArchetypeEntityId) -> Option<Entity> {
        self.ids_to_entity.get(&id).cloned()
    }
}

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

    use rand::Rng;

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

    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, &Tick::default());

        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, &Tick::default());

        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, &Tick::default());
        let earch2 = a.add_entity(e2, b2, &Tick::default());

        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, &Tick::default());
        let earch2 = a.add_entity(e2, b2, &Tick::default());

        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,
                &Tick::default()
            );
        }
        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],),
                &Tick::default()
            );
        }

        // Remove the 'middle' entity in the column
        let moved_entity = a.remove_entity(
            Entity {
                id: EntityId(1u64),
                generation: 0,
            },
            &Tick::default()
        ).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,
            &Tick::default()
        );

        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);
    }
}