From c3c1e819138fcd65ac4ebae66dd8a9f87d4f0093 Mon Sep 17 00:00:00 2001
From: SeanOMik <seanomik@gmail.com>
Date: Tue, 5 Dec 2023 23:22:48 -0500
Subject: [PATCH] Create QueryResourceMut

---
 lyra-ecs/src/query/resource.rs | 107 ++++++++++++++++++++++++++++++---
 1 file changed, 99 insertions(+), 8 deletions(-)

diff --git a/lyra-ecs/src/query/resource.rs b/lyra-ecs/src/query/resource.rs
index a0c19da..7d29b57 100644
--- a/lyra-ecs/src/query/resource.rs
+++ b/lyra-ecs/src/query/resource.rs
@@ -1,4 +1,4 @@
-use std::{marker::PhantomData, any::TypeId, ptr::NonNull, cell::Ref};
+use std::{marker::PhantomData, any::TypeId, ptr::NonNull, cell::{Ref, RefMut}};
 
 use crate::{world::World, resource::ResourceObject};
 
@@ -30,13 +30,11 @@ impl<'a, T: 'a + 'static> Fetch<'a> for FetchResource<'a, T> {
 
 }
 
-/// A query type for getting Resources
-/// 
-/// Resources are stored in an archetype with the id of [`world::RESOURCE_ARCHETYPE_ID`]. There is only one instance of a type of resource.
-/// The resources are stored in that archetype in the first entity in the column.
+/// 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> {}
 
@@ -46,8 +44,6 @@ impl<T: ResourceObject> Clone for QueryResource<T> {
     }
 }
 
-pub type Resource<T> = QueryResource<T>;
-
 impl<T: ResourceObject + 'static> Query for QueryResource<T> {
     type Item<'a> = Ref<'a, T>;
 
@@ -80,9 +76,83 @@ impl<T: ResourceObject + 'static> Query for QueryResource<T> {
 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, _arch_id: crate::archetype::ArchetypeId, _archetype: &'a crate::archetype::Archetype) -> Self::Fetch<'a> {
+        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}};
+    use crate::{world::World, tests::{Vec2, Vec3}, QueryResourceMut};
 
     use super::QueryResource;
 
@@ -129,4 +199,25 @@ mod tests {
         let i = world.view::<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::<QueryResourceMut<SomeCounter>>();
+            let mut resmut = resmut_iter.next().unwrap();
+            assert_eq!(resmut.0, 0);
+            resmut.0 += 20;
+        }
+
+        let mut res_iter = world.view::<QueryResource<SomeCounter>>();
+        let res = res_iter.next().unwrap();
+        assert_eq!(res.0, 20);
+    }
 }
\ No newline at end of file