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

use std::marker::PhantomData;

use atomic_refcell::{AtomicRef, AtomicRefMut};

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

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

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

impl<'a, T: ResourceObject + 'a> Fetch<'a> for FetchResource<'a, T> {
    type Item = Res<'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();
        Res(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> = Res<'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>;
}

/// A struct used for querying resources from the World.
pub struct Res<'a, T: ResourceObject>(pub(crate) AtomicRef<'a, T>);

impl<'a, T: ResourceObject> std::ops::Deref for Res<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.0.deref()
    }
}

impl<'a, T: ResourceObject> AsQuery for Res<'a, T> {
    type Query = QueryResource<T>;
}

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

impl<'a, T: ResourceObject + 'a> Fetch<'a> for FetchResourceMut<'a, T> {
    type Item = ResMut<'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();
        ResMut(w.get_resource_mut::<T>())
    }

}

/// A query type for borrowing Resources mutably.
pub struct QueryResourceMut<T: ResourceObject> {
    _phantom: PhantomData<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> = ResMut<'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>;
}

/// A struct used for querying resources from the World.
pub struct ResMut<'a, T: ResourceObject>(pub(crate) AtomicRefMut<'a, T>);

impl<'a, T: ResourceObject> std::ops::Deref for ResMut<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.0.deref()
    }
}

impl<'a, T: ResourceObject> std::ops::DerefMut for ResMut<'a, T> {

    fn deref_mut(&mut self) -> &mut Self::Target {
        self.0.deref_mut()
    }
}

impl<'a, T: ResourceObject> AsQuery for ResMut<'a, T> {
    type Query = QueryResourceMut<T>;
}

#[cfg(test)]
mod tests {
    use std::ops::{Deref, DerefMut};

    use crate::{query::{Res, ResMut}, tests::{Vec2, Vec3}, World};

    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::<Res<SomeCounter>>();
        let res = res_iter.next().unwrap();
        let res = res.deref();
        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::<(Res<SomeCounter>, &Vec2)>();
        assert_eq!(i.count(), 3);
        let i = world.view_iter::<(&Vec2, Res<SomeCounter>)>();
        assert_eq!(i.count(), 3);

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

        let i = world.view_iter::<Res<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::<ResMut<SomeCounter>>();
            let mut resmut = resmut_iter.next().unwrap();
            let resmut = resmut.deref_mut();
            assert_eq!(resmut.0, 0);
            resmut.0 += 20;
        }

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