From ddc6a3dbf2ae7e4ccec21618a2f94711e692e2fa Mon Sep 17 00:00:00 2001
From: SeanOMik <seanomik@gmail.com>
Date: Fri, 8 Dec 2023 17:35:39 -0500
Subject: [PATCH] Add ArgFetcher for resources

---
 lyra-ecs/src/system/mod.rs | 80 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 77 insertions(+), 3 deletions(-)

diff --git a/lyra-ecs/src/system/mod.rs b/lyra-ecs/src/system/mod.rs
index 9b6d035..800210f 100644
--- a/lyra-ecs/src/system/mod.rs
+++ b/lyra-ecs/src/system/mod.rs
@@ -1,6 +1,6 @@
-use std::{ptr::NonNull, marker::PhantomData};
+use std::{ptr::NonNull, marker::PhantomData, cell::{Ref, RefMut}};
 
-use crate::{world::World, View, Query, Access};
+use crate::{world::World, View, Query, Access, ResourceObject};
 
 pub mod graph;
 
@@ -151,9 +151,55 @@ impl FnArgFetcher for WorldMutArgFetcher {
     }
 }
 
+pub struct ResourceArgFetcher<R: ResourceObject> {
+    phantom: PhantomData<fn() -> R>
+}
+
+impl<'a, R: ResourceObject> FnArg for Ref<'a, R> {
+    type Fetcher = ResourceArgFetcher<R>;
+}
+
+impl<R: ResourceObject> FnArgFetcher for ResourceArgFetcher<R> {
+    type Arg<'a> = Ref<'a, R>;
+
+    fn new() -> Self {
+        ResourceArgFetcher {
+            phantom: PhantomData
+        }
+    }
+
+    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a> {
+        let world = world.as_ref();
+        world.get_resource::<R>()
+    }
+}
+
+pub struct ResourceMutArgFetcher<R: ResourceObject> {
+    phantom: PhantomData<fn() -> R>
+}
+
+impl<'a, R: ResourceObject> FnArg for RefMut<'a, R> {
+    type Fetcher = ResourceMutArgFetcher<R>;
+}
+
+impl<R: ResourceObject> FnArgFetcher for ResourceMutArgFetcher<R> {
+    type Arg<'a> = RefMut<'a, R>;
+
+    fn new() -> Self {
+        ResourceMutArgFetcher {
+            phantom: PhantomData
+        }
+    }
+
+    unsafe fn get<'a>(&mut self, world: NonNull<World>) -> Self::Arg<'a> {
+        let world = world.as_ref();
+        world.get_resource_mut::<R>()
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use std::ptr::NonNull;
+    use std::{ptr::NonNull, cell::RefMut};
 
     use crate::{tests::{Vec2, Vec3}, View, QueryBorrow, world::World};
     use super::{System, IntoSystem};
@@ -230,6 +276,34 @@ mod tests {
 
         test_system.into_system().execute(NonNull::from(&world)).unwrap();
 
+        let test_system = |world: &mut World| -> anyhow::Result<()> {
+            let mut counter = world.get_resource_mut::<SomeCounter>();
+            counter.0 += 10;
+
+            Ok(())
+        };
+
+        test_system.into_system().execute(NonNull::from(&world)).unwrap();
+
+        let counter = world.get_resource::<SomeCounter>();
+        assert_eq!(counter.0, 20);
+    }
+
+    #[test]
+    fn resource_system() {
+        let mut world = World::new();
+        world.spawn((Vec2::rand(), Vec3::rand()));
+        world.spawn((Vec2::rand(), Vec3::rand()));
+        world.add_resource(SomeCounter(0));
+
+        let test_system = |mut counter: RefMut<SomeCounter>| -> anyhow::Result<()> {
+            counter.0 += 10;
+
+            Ok(())
+        };
+
+        test_system.into_system().execute(NonNull::from(&world)).unwrap();
+
         let counter = world.get_resource::<SomeCounter>();
         assert_eq!(counter.0, 10);
     }