From 8dcaae99e177c4ed902569a4f7ed0b0df3b97d2b Mon Sep 17 00:00:00 2001
From: JMS55 <47158642+JMS55@users.noreply.github.com>
Date: Tue, 12 Dec 2023 18:24:48 -0800
Subject: [PATCH] Make use_system() hook return a closure

---
 src/deferred_system.rs | 13 +++++++++----
 src/hooks.rs           |  8 +++-----
 src/lib.rs             |  1 -
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/deferred_system.rs b/src/deferred_system.rs
index c02e293..9ed6ae7 100644
--- a/src/deferred_system.rs
+++ b/src/deferred_system.rs
@@ -9,18 +9,19 @@ use std::sync::Arc;
 
 #[derive(Resource, Default)]
 pub struct DeferredSystemRegistry {
+    // TODO: Replace ref_counts with Box<Vec<SystemId>> and insert SystemId into it on unmount
     pub ref_counts: HashMap<SystemId, Arc<()>>,
     pub run_queue: Box<Vec<SystemId>>,
 }
 
 #[derive(Clone, Copy)]
-pub struct DeferredSystem {
+struct DeferredSystem {
     id: SystemId,
     run_queue: *mut Vec<SystemId>,
 }
 
 impl DeferredSystem {
-    pub fn schedule(&self) {
+    fn schedule(&self) {
         unsafe { &mut *self.run_queue }.push(self.id);
     }
 }
@@ -28,9 +29,13 @@ impl DeferredSystem {
 unsafe impl Send for DeferredSystem {}
 unsafe impl Sync for DeferredSystem {}
 
-pub fn new_deferred_system<S, M>(system: S, world: &mut World) -> (DeferredSystem, Arc<()>)
+pub fn new_deferred_system<S, M>(
+    system: S,
+    world: &mut World,
+) -> (impl Fn() + Send + Sync + Copy, Arc<()>)
 where
     S: IntoSystem<(), (), M> + 'static,
+    M: 'static,
 {
     let id = world.register_system(system);
     let ref_count = Arc::new(());
@@ -45,5 +50,5 @@ where
         run_queue: Box::as_mut(&mut system_registry.run_queue),
     };
 
-    (deferred_system, ref_count)
+    (move || deferred_system.schedule(), ref_count)
 }
diff --git a/src/hooks.rs b/src/hooks.rs
index 2f7e38c..30d6f56 100644
--- a/src/hooks.rs
+++ b/src/hooks.rs
@@ -1,7 +1,4 @@
-use crate::{
-    deferred_system::{new_deferred_system, DeferredSystem},
-    tick::EcsContext,
-};
+use crate::{deferred_system::new_deferred_system, tick::EcsContext};
 use bevy::{
     ecs::{
         component::ComponentId,
@@ -93,9 +90,10 @@ where
     }
 }
 
-pub fn use_system<S, M>(cx: &ScopeState, system: S) -> DeferredSystem
+pub fn use_system<S, M>(cx: &ScopeState, system: S) -> impl Fn() + Send + Sync + Copy
 where
     S: IntoSystem<(), (), M> + 'static,
+    M: 'static,
 {
     cx.use_hook(|| new_deferred_system(system, EcsContext::get_world(cx)))
         .0
diff --git a/src/lib.rs b/src/lib.rs
index 1c953e0..314786a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -19,7 +19,6 @@ use bevy::{
 };
 use dioxus::core::{Element, ElementId, Scope};
 
-pub use self::deferred_system::DeferredSystem;
 pub use bevy_mod_picking;
 pub use dioxus;