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

use std::{any::TypeId, ptr::{NonNull, self}, alloc::{self, Layout, alloc, dealloc}, mem, collections::HashMap};

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

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

#[allow(dead_code)]
impl ComponentColumn {
    /// Creates an invalid component column. Do not attempt to use it.
    pub unsafe fn dangling() -> Self {
        ComponentColumn {
            data: NonNull::dangling(),
            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, capacity);
        
        let size = info.layout.size();
        Self {
            data,
            capacity,
            info,
            entry_size: size,
            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 dest = NonNull::new_unchecked(self.data.as_ptr().add(entity_index * self.entry_size));
        ptr::copy_nonoverlapping(comp_src.as_ptr(), dest.as_ptr(), self.entry_size);
        self.len += 1;
    }

    /// 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 ptr = NonNull::new_unchecked(self.data.as_ptr()
            .add(entity_index * self.entry_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 p = self.data.as_ptr()
            .cast::<T>()
            .add(entity_index * self.entry_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 new_ptr = Self::alloc(self.info.layout, new_capacity);
        ptr::copy_nonoverlapping(self.data.as_ptr(), new_ptr.as_ptr(), self.capacity * self.entry_size);

        let old_layout = Layout::from_size_align_unchecked(
            self.info.layout.size().checked_mul(self.capacity).unwrap(),
            self.info.layout.align()
        );

        mem::swap(&mut self.data, &mut new_ptr); // 'new_ptr' is now the old pointer
        dealloc(new_ptr.as_ptr(), old_layout);
        
        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 old_comp_ptr = NonNull::new_unchecked(self.data.as_ptr()
            .add(entity_index * self.entry_size));

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

            ptr::copy_nonoverlapping(new_comp_ptr.as_ptr(), old_comp_ptr.as_ptr(), self.entry_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
    }
}

#[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(crate) id: ArchetypeId,
    pub(crate) entities: HashMap<Entity, ArchetypeEntityId>,
    pub(crate) columns: Vec<ComponentColumn>,
    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.columns.iter_mut().find(|c| c.info.type_id == type_id).unwrap();
            unsafe { col.set_at(entity_index, data); }
        });
        
        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(|(e, 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());
            }
        }
        
        // 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(crate) fn is_archetype_for(&self, types: Vec<TypeId>) -> bool {
        self.columns.iter().all(|c| types.contains(&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); }
        }
    }
}

#[cfg(test)]
mod tests {
    use rand::Rng;

    use crate::{tests::{Vec2, Vec3}, world::{Entity, EntityId}, bundle::{Bundle, self}};

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

    #[test]
    fn column_growth() {
        let mut rng = rand::thread_rng();
        let bundle_count = rng.gen_range(50..150);
        
        let mut bundles = 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::new(rng.gen_range(10.0..3000.0), rng.gen_range(10.0..3000.0)),);
            bundles.push(c);

            a.add_entity(
                Entity {
                    id: EntityId(i as u64),
                    generation: 0
                },
                c
            );
        }

        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 mut rng = rand::thread_rng();
        let range = 30.0..1853.0;
        
        let bundles = vec![
            ( Vec2::new(rng.gen_range(range.clone()), rng.gen_range(range.clone())), ),
            ( Vec2::new(rng.gen_range(range.clone()), rng.gen_range(range.clone())), ),
            ( Vec2::new(rng.gen_range(range.clone()), rng.gen_range(range.clone())), )
        ];

        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].0);
    }
}