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

use std::{any::{Any, TypeId, type_name}, ptr::{NonNull, self}, alloc::{self, Layout, alloc}, mem::size_of};

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

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

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

    /// Creates an empty column of the same type
    pub unsafe fn create_empty(&self, capacity: usize) -> Self {
        let data = ComponentColumn::alloc(self.info.layout, capacity);

        let size = self.info.layout.size();
        Self {
            data,
            capacity,
            info: self.info.clone(),
            entry_size: size,
        }
    }

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

    /// 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 p = self.data.as_ptr()
            .cast::<T>()
            .add(entity_index * self.entry_size);
        &*p
    }

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

#[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: Vec<Entity>,
    pub(crate) columns: Vec<ComponentColumn>,
}

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

impl Archetype {
    /// Create a new archetype from another archetype and add a column
    pub fn new_archetype_add<T: Component>(new_id: ArchetypeId, archetype: &Archetype) -> Archetype {
        let mut columns: Vec<_> = archetype
            .columns
            .iter()
            .map(|c| unsafe { c.create_empty(DEFAULT_CAPACITY) })
            .collect();

        // Make sure a column for the new component does not exist
        assert_ne!(true, columns.iter().any(|c| c.info.type_id == TypeId::of::<T>()) );
        columns.push(unsafe { ComponentColumn::new(ComponentInfo::new::<T>(), DEFAULT_CAPACITY) });

        Archetype {
            id: new_id,
            entities: Vec::new(),
            columns,
        }
    }

    /// Create a new archetype from another archetype and remove a column
    pub fn new_archetype_remove<T: Component>(new_id: ArchetypeId, archetype: &Archetype) -> Archetype {
        let mut columns: Vec<_> = archetype
            .columns
            .iter()
            .map(|c| unsafe { c.create_empty(DEFAULT_CAPACITY) })
            .collect();

        let idx = columns
            .iter()
            .position(|column| column.info.type_id == TypeId::of::<T>())
            .unwrap();
        columns.remove(idx);

        Archetype {
            id: new_id,
            entities: Vec::new(),
            columns,
        }
    }

    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: Vec::new(),
            columns,
        }
    }

    pub fn get_component_mut<T: Component>(&mut self, entity: ArchetypeEntityId) -> Option<&mut T> {
        todo!()
    }

    pub fn get_entity_component<T: Component>(&self, entity: ArchetypeEntityId) -> Option<&T> {
        let type_id = TypeId::of::<T>();
        
        self.columns.iter().find(|c| c.info.type_id == type_id)
            .map(|c| unsafe { c.get(entity.0 as usize) })
    }

    pub fn get_component_column<T: Component>(&self) -> Option<&Vec<T>> {
        todo!()
    }

    /// 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
    {
        let entity_index = self.entities.len();
        self.entities.push(entity);

        bundle.takefn(|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)
    }

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