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

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

use crate::{archetype::{ArchetypeId, Archetype}, bundle::Bundle, component::Component, query::{ViewIter, View, DefaultQuery}, resource::ResourceData};

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

pub struct Record {
    id: ArchetypeId,
    index: ArchetypeEntityId,
}

pub struct World {
    pub(crate) archetypes: HashMap<ArchetypeId, Archetype>,
    next_archetype_id: ArchetypeId,
    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.clone()));

        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 view<'a, T: 'static + Component + DefaultQuery>(&'a self) -> ViewIter<'a, T::Query> {
        let archetypes = self.archetypes.values().collect();
        let v = View::new(self, T::default_query(), archetypes);
        v.into_iter()
    }

    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;

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