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

use std::{collections::{HashMap, VecDeque}, any::TypeId, cell::{Ref, RefMut}, ptr::NonNull};

use crate::{archetype::{ArchetypeId, Archetype}, bundle::Bundle, component::Component, query::{ViewIter, View}, resource::ResourceData, Query, AsQuery, dynamic::{DynamicViewIter, QueryDynamicType, DynamicView}, ViewOne, ComponentInfo, DynTypeId};

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

/// The id of the entity for the Archetype.
/// The Archetype struct uses this as the index in the component columns
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct ArchetypeEntityId(pub u64);

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Entity {
    pub(crate) id: EntityId,
    pub(crate) generation: u64,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Record {
    pub id: ArchetypeId,
    pub index: ArchetypeEntityId,
}

pub struct World {
    pub(crate) archetypes: HashMap<ArchetypeId, Archetype>,
    next_archetype_id: ArchetypeId,
    pub(crate) entity_index: HashMap<EntityId, Record>,
    dead_entities: VecDeque<Entity>,
    next_entity_id: EntityId,
    resources: HashMap<TypeId, ResourceData>,
}

impl World {
    pub fn new() -> Self {
        Self {
            archetypes: HashMap::new(),
            next_archetype_id: ArchetypeId(0),
            entity_index: HashMap::new(),
            dead_entities: VecDeque::new(),
            next_entity_id: EntityId(0),
            resources: HashMap::new(),
        }
    }

    fn get_new_entity(&mut self) -> Entity {
        match self.dead_entities.pop_front() {
            Some(e) => e,
            None => {
                let new_id = self.next_entity_id;
                self.next_entity_id.0 += 1;

                Entity {
                    id: new_id,
                    generation: 0,
                }
            }
        }
    }

    pub fn spawn<B>(&mut self, bundle: B) -> Entity
    where
        B: Bundle
    {
        let bundle_types = bundle.type_ids();
        let new_entity = self.get_new_entity();

        // try to find an archetype
        let archetype = self.archetypes
            .values_mut()
            .find(|a| a.is_archetype_for(&bundle_types));

        if let Some(archetype) = archetype {
            let arche_idx = archetype.add_entity(new_entity, bundle);

            // Create entity record and store it
            let record = Record {
                id: archetype.id,
                index: arche_idx,
            };
            
            self.entity_index.insert(new_entity.id, record);
        }
        // create a new archetype if one isn't found
        else {
            // create archetype
            let new_arch_id = self.next_archetype_id.increment();
            let mut archetype = Archetype::from_bundle_info(new_arch_id, bundle.info());
            let entity_arch_id = archetype.add_entity(new_entity, bundle);

            // store archetype
            self.archetypes.insert(new_arch_id, archetype);

            // Create entity record and store it
            let record = Record {
                id: new_arch_id,
                // this is the first entity in the archetype
                index: entity_arch_id,
            };

            self.entity_index.insert(new_entity.id, record);
        }
        
        new_entity
    }

    /// Despawn an entity from the World
    pub fn despawn(&mut self, entity: Entity) {
        if let Some(record) = self.entity_index.get_mut(&entity.id) {
            let arch = self.archetypes.get_mut(&record.id).unwrap();

            if let Some((moved, new_index)) = arch.remove_entity(entity) {
                // replace the archetype index of the moved index with its new index.
                self.entity_index.get_mut(&moved.id).unwrap().index = new_index;
            }
        }
    }

    pub fn insert<B>(&mut self, entity: Entity, bundle: B)
    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.

        let record = *self.entity_index.get(&entity.id).unwrap();
        let current_arch = self.archetypes.get(&record.id).unwrap();
        
        let mut col_types: Vec<DynTypeId> = current_arch.columns.iter().map(|c| c.info.type_id).collect();
        let orig_col = col_types.clone();
        col_types.extend(bundle.type_ids().into_iter());
        
        let mut col_infos: Vec<ComponentInfo> = current_arch.columns.iter().map(|c| c.info).collect();
        col_infos.extend(bundle.info().into_iter());

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

            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));
                    let col = arch.get_column_mut(col_type).unwrap();
                    col.set_at(res_index.0 as _, ptr);
                }
            }

            bundle.take(|data, type_id, _size| {
                let col = arch.get_column_mut(type_id).unwrap();
                unsafe { col.set_at(res_index.0 as _, data); }
                col.len += 1;
            });

            arch.entities.insert(entity, res_index);
            
            let new_record = Record {
                id: arch.id,
                index: res_index,
            };
            self.entity_index.insert(entity.id, new_record);
        } else {
            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);

            self.archetypes.insert(new_arch_id, archetype);

            // Create entity record and store it
            let record = Record {
                id: new_arch_id,
                index: entity_arch_id,
            };

            self.entity_index.insert(entity.id, record);
        }

        let current_arch = self.archetypes.get_mut(&record.id).unwrap();
        current_arch.remove_entity(entity);
    }

    /// View into the world for a set of entities that satisfy the queries.
    pub fn view_iter<'a, T: 'static + Component + AsQuery>(&'a self) -> ViewIter<'a, T::Query> {
        let archetypes = self.archetypes.values().collect();
        let v = View::new(self, T::Query::new(), archetypes);
        v.into_iter()
    }

    pub fn dynamic_view(&self) -> DynamicView {
        DynamicView::new(self)
    }

    pub fn view_one<'a, T: 'static + AsQuery>(&'a self, entity: Entity) -> ViewOne<'a, T::Query> {
        ViewOne::new(self, entity.id, T::Query::new())
    }

    //pub fn view_one(&self, entity: EntityId) -> 

    pub fn add_resource<T: 'static>(&mut self, data: T) {
        self.resources.insert(TypeId::of::<T>(), ResourceData::new(data));
    }

    /// Gets a resource from the World.
    /// 
    /// Will panic if the resource is not in the world. See [`try_get_resource`] for
    /// a function that returns an option.
    pub fn get_resource<'a, T: 'static>(&'a self) -> Ref<'a, T> {
        self.resources.get(&TypeId::of::<T>()).unwrap()
            .get()
    }

    /// Attempts to get a resource from the World.
    /// 
    /// Returns `None` if the resource was not found.
    pub fn try_get_resource<'a, T: 'static>(&'a self) -> Option<Ref<'a, T>> {
        self.resources.get(&TypeId::of::<T>())
            .map(|r| r.try_get())
            .flatten()
    }

    /// Gets a mutable borrow of a resource from the World.
    /// 
    /// Will panic if the resource is not in the world. See [`try_get_resource_mut`] for
    /// a function that returns an option.
    pub fn get_resource_mut<'a, T: 'static>(&'a self) -> RefMut<'a, T> {
        self.resources.get(&TypeId::of::<T>()).unwrap()
            .get_mut()
    }

    /// Attempts to get a mutable borrow of a resource from the World.
    /// 
    /// Returns `None` if the resource was not found.
    pub fn try_get_resource_mut<'a, T: 'static>(&'a self) -> Option<RefMut<'a, T>> {
        self.resources.get(&TypeId::of::<T>())
            .map(|r| r.try_get_mut())
            .flatten()
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{Vec2, Vec3};

    use super::World;

    struct SimpleCounter(i32);

    #[test]
    fn spawning_entity() {
        let mut world = World::new();
        let _e = world.spawn((Vec2 {
            x: 10.0,
            y: 15.0,
        }, ));
    }

    #[test]
    fn world_view_entities() {
        let mut world = World::new();
        world.spawn((Vec2 {
            x: 10.0,
            y: 15.0,
        }, ));
        world.spawn((Vec2 {
            x: 152.0,
            y: 3585.0,
        }, ));
        world.spawn((Vec2 {
            x: 235.0,
            y: 734.0,
        }, ));

        let mut count = 0;
        for pos in world.view_iter::<&Vec2>() {
            println!("Found entity at {:?}", pos);
            count += 1;
        }
        assert!(count == 3);
    }

    #[test]
    fn despawn_entity() {
        let mut world = World::new();
        world.spawn((Vec2::rand(),));
        let middle_en = world.spawn((Vec2::rand(),));
        let last_en = world.spawn((Vec2::rand(),));

        world.despawn(middle_en);

        let record = world.entity_index.get(&last_en.id).unwrap();
        assert_eq!(record.index.0, 1);
    }

    #[test]
    fn simple_resource() {
        let mut world = World::new();
        {
            let counter = SimpleCounter(0);
            world.add_resource(counter);
        }

        let counter = world.get_resource::<SimpleCounter>();
        assert_eq!(counter.0, 0);
        drop(counter);

        let mut counter = world.get_resource_mut::<SimpleCounter>();
        counter.0 += 4582;
        drop(counter);

        assert!(world.try_get_resource::<u32>().is_none());
    }

    #[test]
    fn resource_multi_borrow() {
        let mut world = World::new();
        {
            let counter = SimpleCounter(4582);
            world.add_resource(counter);
        }

        // test multiple borrows at the same time
        let counter = world.get_resource::<SimpleCounter>();
        assert_eq!(counter.0, 4582);
        let counter2 = world.get_resource::<SimpleCounter>();
        assert_eq!(counter2.0, 4582);
        assert_eq!(counter2.0, 4582);
    }

    #[test]
    fn resource_one_mutable_borrow() {
        let mut world = World::new();
        {
            let counter = SimpleCounter(4582);
            world.add_resource(counter);
        }

        // test multiple borrows at the same time
        let counter = world.get_resource_mut::<SimpleCounter>();
        assert_eq!(counter.0, 4582);
        assert!(world.try_get_resource_mut::<SimpleCounter>().is_none());
        assert_eq!(counter.0, 4582);
    }

    #[test]
    fn insert_into_existing_archetype() {
        let mut world = World::new();
        let e = world.spawn((Vec2::rand(),));
        world.spawn((Vec2::rand(),Vec3::rand()));

        world.insert(e, (Vec3::rand(),));

        assert!(world.view_one::<&Vec3>(e).get().is_some())
    }

    #[test]
    fn insert_into_new_archetype() {
        let mut world = World::new();
        let e = world.spawn((Vec2::rand(),));

        world.insert(e, (Vec3::rand(),));

        assert!(world.view_one::<&Vec3>(e).get().is_some())
    }

    #[test]
    fn view_one() {
        let v = Vec2::rand();

        let mut world = World::new();
        let e = world.spawn((v,));

        let view = world.view_one::<&Vec2>(e);
        assert_eq!(*view.get().unwrap(), v);
    }
}