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

use std::{any::TypeId, collections::HashMap, ptr::NonNull};

use atomic_refcell::{AtomicRef, AtomicRefMut};

use crate::{archetype::{Archetype, ArchetypeId}, bundle::Bundle, query::{dynamic::DynamicView, AsQuery, Query, ViewState, ViewIter, ViewOne}, resource::ResourceData, ComponentInfo, DynTypeId, Entities, Entity, ResourceObject, Tick, TickTracker};

/// The id of the entity for the Archetype.
/// 
/// The Archetype 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 Record {
    pub id: ArchetypeId,
    pub index: ArchetypeEntityId,
}

pub struct World {
    pub(crate) archetypes: HashMap<ArchetypeId, Archetype>,
    next_archetype_id: ArchetypeId,
    resources: HashMap<TypeId, ResourceData>,
    tracker: TickTracker,
    pub(crate) entities: Entities,
}

impl Default for World {
    fn default() -> Self {
        Self {
            archetypes: HashMap::new(),
            next_archetype_id: ArchetypeId(0),
            resources: HashMap::new(),
            tracker: TickTracker::new(),
            entities: Entities::default(),
        }
    }
}

impl World {
    pub fn new() -> Self {
        Self::default()
    }

    /// Reserves an entity in the world
    pub fn reserve_entity(&mut self) -> Entity {
        self.entities.reserve()
    }

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

    /// Spawn the components into a reserved entity. Only do this with entities that
    /// were reserved with [`World::reserve_entity`].
    ///
    /// # Safety
    /// Do not use this method with an entity that is currently alive, it WILL cause undefined behavior.
    pub fn spawn_into<B>(&mut self, entity: Entity, bundle: B)
    where
        B: Bundle
    {
        let bundle_types = bundle.type_ids();

        let tick = self.tick();

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

        if let Some(archetype) = archetype {
            // make at just one check to ensure you're not spawning twice
            debug_assert!(!archetype.entity_ids.contains_key(&entity),
                "You attempted to spawn components into an entity that already exists!");

            let arche_idx = archetype.add_entity(entity, bundle, &tick);

            // Create entity record and store it
            let record = Record {
                id: archetype.id(),
                index: arche_idx,
            };
            
            self.entities.insert_entity_record(entity, 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(entity, bundle, &tick);

            // 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.entities.insert_entity_record(entity, record);
        }
    }

    /// Despawn an entity from the World
    pub fn despawn(&mut self, entity: Entity) {
        // Tick the tracker if the entity is spawned. This is done here instead of the `if let`
        // below due to the borrow checker complaining about multiple mutable borrows to self.
        let tick = if self.entities.arch_index.contains_key(&entity.id) {
            Some(self.tick())
        } else { None };

        if let Some(record) = self.entities.arch_index.get_mut(&entity.id) {
            let tick = tick.unwrap();
            let arch = self.archetypes.get_mut(&record.id).unwrap();

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

    /// Insert a bundle into an existing entity. If the components are already existing on the
    /// entity, they will be updated, else the entity will be moved to a different Archetype
    /// that can store the entity. That may involve creating a new Archetype.
    pub fn insert<B>(&mut self, entity: Entity, bundle: B)
    where
        B: Bundle
    {
        // TODO: If the entity already has the components in `bundle`, update the values of the
        // components with the bundle.

        let tick = self.tick();

        let record = self.entities.entity_record(entity).unwrap();
        let current_arch = self.archetypes.get(&record.id).unwrap();
        let current_arch_len = current_arch.len();

        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());
        
        let mut col_infos: Vec<ComponentInfo> = current_arch.columns.iter().map(|c| c.info).collect();
        col_infos.extend(bundle.info());

        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(record.index.0 as usize * col_info.layout().size()));
                    let col = arch.get_column_mut(col_type).unwrap();
                    // set_at is used since the entity was reserved
                    col.set_at(res_index.0 as _, ptr, tick);
                }
            }

            bundle.take(|data, type_id, _size| {
                let col = arch.get_column_mut(type_id).unwrap();
                // set_at is used since the entity was reserved
                unsafe { col.set_at(res_index.0 as _, data, tick); }
            });

            arch.entity_ids.insert(entity, res_index);
            arch.ensure_synced();
            
            let new_record = Record {
                id: arch.id(),
                index: res_index,
            };
            self.entities.insert_entity_record(entity, new_record);
        } else {
            if current_arch_len == 1 {
                // if this entity is the only entity for this archetype, add more columns to it
                let current_arch = self.archetypes.get_mut(&record.id).unwrap();
                current_arch.extend(&tick, vec![bundle]);

                return;
            }

            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, &tick);
            
            // move the old components into the new archetype
            for (column_ptr, column_info) in col_ptrs.into_iter() {
                unsafe {
                    // ptr of component for the entity
                    let comp_ptr = NonNull::new_unchecked(column_ptr.as_ptr()
                        .add(record.index.0 as usize * column_info.layout().size()));
                    
                    let col = archetype.get_column_mut(column_info.type_id()).unwrap();
                    col.insert_entity(entity_arch_id.0 as _, comp_ptr, tick);
                }
            }

            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.entities.insert_entity_record(entity, record);
        }

        let current_arch = self.archetypes.get_mut(&record.id).unwrap();
        if current_arch.len() > 1 {
            current_arch.remove_entity(entity, &tick);
        } else if current_arch.len() == 1 {
            // The old archetype will only be removed if there was another archetype that would
            // work for the entity's components, and the old archetype only had a single entity.
            self.archetypes.remove(&record.id).unwrap();
        }
    }

    pub fn entity_archetype(&self, entity: Entity) -> Option<&Archetype> {
        self.entities.entity_record(entity)
            .and_then(|record| self.archetypes.get(&record.id))
    }

    pub fn entity_archetype_mut(&mut self, entity: Entity) -> Option<&mut Archetype> {
        self.entities.entity_record(entity)
            .and_then(|record| self.archetypes.get_mut(&record.id))
    }

    /// View into the world for a set of entities that satisfy the queries.
    pub fn view<Q: AsQuery>(&self) -> ViewState<Q::Query, ()> {
        self.filtered_view::<Q, ()>()
    }

    /// View into the world for a set of entities that satisfy the query and the filter.
    pub fn filtered_view<Q: AsQuery, F: AsQuery>(&self) -> ViewState<Q::Query, F::Query> {
        let archetypes = self.archetypes.values().collect();
        ViewState::<Q::Query, F::Query>::new(self, Q::Query::new(), F::Query::new(), archetypes)
    }

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

    /// View into the world for a set of entities that satisfy the queries.
    pub fn filtered_view_iter<Q: AsQuery, F: AsQuery>(&self) -> ViewIter<Q::Query, F::Query> {
        let archetypes = self.archetypes.values().collect();
        let v = ViewState::new(self, Q::Query::new(), F::Query::new(), archetypes);
        v.into_iter()
    }

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

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

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

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

    pub fn add_resource_default<T: ResourceObject + Default>(&mut self) {
        self.resources.insert(TypeId::of::<T>(), ResourceData::new(T::default()));
    }

    /// Get a resource from the world, or insert it into the world with the provided
    /// `fn` and return it.
    pub fn get_resource_or_else<T: ResourceObject, F>(&mut self, f: F) -> AtomicRefMut<T>
    where
        F: Fn() -> T + 'static
    {
        self.resources.entry(TypeId::of::<T>())
            .or_insert_with(|| ResourceData::new(f()))
            .get_mut()
    }

    /// Get a resource from the world, or insert it into the world as its default.
    pub fn get_resource_or_default<T: ResourceObject + Default>(&mut self) -> AtomicRefMut<T>
    {
        self.resources.entry(TypeId::of::<T>())
            .or_insert_with(|| ResourceData::new(T::default()))
            .get_mut()
    }

    /// Gets a resource from the World.
    /// 
    /// Will panic if the resource is not in the world. See [`World::try_get_resource`] for
    /// a function that returns an option.
    pub fn get_resource<T: ResourceObject>(&self) -> AtomicRef<T> {
        self.resources.get(&TypeId::of::<T>())
            .expect(&format!("World is missing resource of type '{}'", std::any::type_name::<T>()))
            .get()
    }

    /// Returns boolean indicating if the World contains a resource of type `T`.
    pub fn has_resource<T: ResourceObject>(&self) -> bool {
        self.resources.contains_key(&TypeId::of::<T>())
    }

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

    /// Gets a mutable borrow of a resource from the World.
    /// 
    /// Will panic if the resource is not in the world. See [`World::try_get_resource_mut`] for
    /// a function that returns an option.
    pub fn get_resource_mut<T: ResourceObject>(&self) -> AtomicRefMut<T> {
        self.resources.get(&TypeId::of::<T>())
            .expect(&format!("World is missing resource of type '{}'", std::any::type_name::<T>()))
            .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<T: ResourceObject>(&self) -> Option<AtomicRefMut<T>> {
        self.resources.get(&TypeId::of::<T>())
            .and_then(|r| r.try_get_mut())
    }

    /// Increments the TickTracker which is used for tracking changes to components.
    /// 
    /// Most users wont need to call this manually, its done for you through queries and views.
    pub fn tick(&self) -> Tick {
        self.tracker.tick()
    }

    /// Gets the current tick that the world is at.
    /// 
    /// See [`TickTracker`]
    pub fn current_tick(&self) -> Tick {
        self.tracker.current()
    }

    pub fn tick_tracker(&self) -> &TickTracker {
        &self.tracker
    }

    /// Attempts to find a resource in the world and returns a NonNull pointer to it
    pub unsafe fn try_get_resource_ptr<T: ResourceObject>(&self) -> Option<NonNull<T>> {
        self.resources.get(&TypeId::of::<T>())
            .map(|d| unsafe { NonNull::new_unchecked(d.data.as_ptr() as *mut T) })
    }
}

// TODO: Ensure that all non-send resources are only accessible on the main thread.
unsafe impl Send for World {}
unsafe impl Sync for World {}

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

    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.entities.entity_record(last_en).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!(counter.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 that its only possible to get a single mutable borrow
        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 insert_multiple_times() {
        let v2s = &[Vec2::rand(), Vec2::rand(), Vec2::rand()];
        let v3s = &[Vec3::rand(), Vec3::rand(), Vec3::rand()];

        let mut world = World::new();
        let e1 = world.spawn(v2s[0]);
        let e2 = world.spawn(v2s[1]);
        let e3 = world.spawn(v2s[2]);
        println!("Spawned entities");

        let ev2 = world.view_one::<&Vec2>(e2).get()
            .expect("Failed to find Vec2 and Vec3 on inserted entity!");
        assert_eq!(*ev2, v2s[1]);
        drop(ev2);

        let insert_and_assert = |world: &mut World, e: Entity, v2: Vec2, v3: Vec3| {
            println!("inserting entity");
            world.insert(e, (v3,));
            println!("inserted entity");

            let (ev2, ev3) = world.view_one::<(&Vec2, &Vec3)>(e).get()
                .expect("Failed to find Vec2 and Vec3 on inserted entity!");
            assert_eq!(*ev2, v2);
            assert_eq!(*ev3, v3);
        };

        insert_and_assert(&mut world, e2, v2s[1], v3s[1]);
        println!("Entity 2 is good");
        insert_and_assert(&mut world, e3, v2s[2], v3s[2]);
        println!("Entity 3 is good");
        assert_eq!(world.archetypes.len(), 2);
        println!("No extra archetypes were created");

        insert_and_assert(&mut world, e1, v2s[0], v3s[0]);
        println!("Entity 1 is good");
        assert_eq!(world.archetypes.len(), 1);
        println!("Empty archetype was removed");
    }

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

    #[test]
    fn view_change_tracking() {
        let mut world = World::new();

        println!("spawning");
        world.spawn((Vec2::new(10.0, 10.0),));
        world.spawn((Vec2::new(5.0, 5.0),));
        println!("spawned");

        for mut v in world.view_iter::<&mut Vec2>() {
            v.y += 50.0;
            println!("Moved v to {:?}", v);
        }

        let world_tick = world.current_tick();
        println!("The world tick is {}", *world_tick);
        for (v, tick) in world.view_iter::<(&Vec2, TickOf<Vec2>)>() {
            println!("Is at {:?}, it was changed at {}", v, *tick);
            assert!(v.y > 50.0);
            assert!(tick >= world_tick);
        }
    }
}