lyra-engine/lyra-ecs/src/query/resource.rs

use std::{marker::PhantomData, cell::{Ref, RefMut}};

use crate::{world::World, resource::ResourceObject};

use super::{Query, Fetch, AsQuery};

pub struct FetchResource<'a, T> {
    world: Option<&'a World>,
    _phantom: PhantomData<T>,
}

impl<'a, T: 'a + 'static> Fetch<'a> for FetchResource<'a, T> {
    type Item = Ref<'a, T>;

    fn dangling() -> Self {
        Self {
            world: None,
            _phantom: PhantomData,
        }
    }

    fn can_visit_item(&mut self, _entity: crate::world::ArchetypeEntityId) -> bool {
        true
    }

    unsafe fn get_item(&mut self, _entity: crate::world::ArchetypeEntityId) -> Self::Item {
        let w = self.world.unwrap();
        w.get_resource::<T>()
    }

}

/// A query type for borrowing Resources.
pub struct QueryResource<T: ResourceObject> {
    _phantom: PhantomData<T>
}
pub type Resource<T> = QueryResource<T>;

impl<T: ResourceObject> Copy for QueryResource<T> {}

impl<T: ResourceObject> Clone for QueryResource<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: ResourceObject + 'static> Query for QueryResource<T> {
    type Item<'a> = Ref<'a, T>;

    type Fetch<'a> = FetchResource<'a, T>;

    const ALWAYS_FETCHES: bool = true;

    fn new() -> Self {
        QueryResource::<T> {
            _phantom: PhantomData
        }
    }

    fn can_visit_archetype(&self, _archetype: &crate::archetype::Archetype) -> bool {
        true
    }

    unsafe fn fetch<'a>(&self, world: &'a World, _archetype: &'a crate::archetype::Archetype, tick: crate::Tick) -> Self::Fetch<'a> {
        let _ = tick;
        self.fetch_world(world).unwrap()
    }

    unsafe fn fetch_world<'a>(&self, world: &'a World) -> Option<Self::Fetch<'a>> {
        Some(FetchResource {
            world: Some(world),
            _phantom: PhantomData,
        })
    }
}

impl<R: ResourceObject> AsQuery for QueryResource<R> {
    type Query = QueryResource<R>;
}

pub struct FetchResourceMut<'a, T> {
    world: Option<&'a World>,
    _phantom: PhantomData<T>,
}

impl<'a, T: 'a + 'static> Fetch<'a> for FetchResourceMut<'a, T> {
    type Item = RefMut<'a, T>;

    fn dangling() -> Self {
        Self {
            world: None,
            _phantom: PhantomData,
        }
    }

    fn can_visit_item(&mut self, _entity: crate::world::ArchetypeEntityId) -> bool {
        true
    }

    unsafe fn get_item(&mut self, _entity: crate::world::ArchetypeEntityId) -> Self::Item {
        let w = self.world.unwrap();
        w.get_resource_mut::<T>()
    }

}

/// A query type for borrowing Resources mutably.
pub struct QueryResourceMut<T: ResourceObject> {
    _phantom: PhantomData<T>
}

pub type ResourceMut<T> = QueryResourceMut<T>;

impl<T: ResourceObject> Copy for QueryResourceMut<T> {}

impl<T: ResourceObject> Clone for QueryResourceMut<T> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<T: ResourceObject + 'static> Query for QueryResourceMut<T> {
    type Item<'a> = RefMut<'a, T>;

    type Fetch<'a> = FetchResourceMut<'a, T>;

    const ALWAYS_FETCHES: bool = true;

    fn new() -> Self {
        QueryResourceMut::<T> {
            _phantom: PhantomData
        }
    }

    fn can_visit_archetype(&self, _archetype: &crate::archetype::Archetype) -> bool {
        true
    }

    unsafe fn fetch<'a>(&self, world: &'a World, _archetype: &'a crate::archetype::Archetype, tick: crate::Tick) -> Self::Fetch<'a> {
        let _ = tick;
        self.fetch_world(world).unwrap()
    }

    unsafe fn fetch_world<'a>(&self, world: &'a World) -> Option<Self::Fetch<'a>> {
        Some(FetchResourceMut {
            world: Some(world),
            _phantom: PhantomData,
        })
    }
}

impl<R: ResourceObject> AsQuery for QueryResourceMut<R> {
    type Query = QueryResourceMut<R>;
}
#[cfg(test)]
mod tests {
    use crate::{world::World, tests::{Vec2, Vec3}, query::QueryResourceMut};

    use super::QueryResource;

    struct SomeCounter(u32);

    #[test]
    fn simple_query() {
        let mut world = World::new();
        {
            let counter = SomeCounter(0);
            world.add_resource(counter);
            println!("Added resource");
        }

        let mut res_iter = world.view_iter::<QueryResource<SomeCounter>>();
        let res = res_iter.next().unwrap();
        assert_eq!(res.0, 0);
    }

    #[test]
    fn complex_query() {
        let mut world = World::new();
        {
            let counter = SomeCounter(0);
            
            world.spawn((Vec2::rand(),));
            world.spawn((Vec2::rand(),));
            world.spawn((Vec2::rand(),));
            world.spawn((Vec3::rand(),));

            world.add_resource(counter);
            println!("Added resource");
        }

        let i = world.view_iter::<(QueryResource<SomeCounter>, &Vec2)>();
        assert_eq!(i.count(), 3);
        let i = world.view_iter::<(&Vec2, QueryResource<SomeCounter>)>();
        assert_eq!(i.count(), 3);

        for (res, e) in world.view_iter::<(QueryResource<SomeCounter>, &Vec2)>() {
            println!("Got res {}! and entity at {:?}", res.0, e);
        }

        let i = world.view_iter::<QueryResource<SomeCounter>>();
        assert_eq!(i.count(), 1);
    }

    #[test]
    fn mutate_query() {
        let mut world = World::new();
        {
            let counter = SomeCounter(0);
            world.add_resource(counter);
            println!("Added resource");
        }

        {
            let mut resmut_iter = world.view_iter::<QueryResourceMut<SomeCounter>>();
            let mut resmut = resmut_iter.next().unwrap();
            assert_eq!(resmut.0, 0);
            resmut.0 += 20;
        }

        let mut res_iter = world.view_iter::<QueryResource<SomeCounter>>();
        let res = res_iter.next().unwrap();
        assert_eq!(res.0, 20);
    }
}