From 3232fb392ab80057e15b198434b0f0165fc852c6 Mon Sep 17 00:00:00 2001
From: JMS55 <47158642+JMS55@users.noreply.github.com>
Date: Thu, 7 Dec 2023 12:04:06 -0800
Subject: [PATCH] Fix DeferredSystem unsafety

---
 src/deferred_system.rs | 27 ++++++---------------------
 src/hooks.rs           | 10 ++--------
 src/tick.rs            |  2 +-
 3 files changed, 9 insertions(+), 30 deletions(-)

diff --git a/src/deferred_system.rs b/src/deferred_system.rs
index bb734e5..533870d 100644
--- a/src/deferred_system.rs
+++ b/src/deferred_system.rs
@@ -6,7 +6,7 @@ use bevy::ecs::{
 #[derive(Clone, Copy)]
 pub struct DeferredSystem {
     id: SystemId,
-    world: *mut World,
+    run_queue: *mut Vec<SystemId>,
 }
 
 impl DeferredSystem {
@@ -15,33 +15,18 @@ impl DeferredSystem {
         S: IntoSystem<(), (), ()> + 'static,
     {
         Self {
-            id: world.register_system(system),
-            world,
+            id: world.register_system(system), // TODO: We never unregister the system
+            run_queue: Box::as_mut(&mut world.resource_mut::<DeferredSystemRunQueue>().0),
         }
     }
 
     pub fn schedule(&self) {
-        // TODO: This is not sound. Pointer to world won't be valid across frames.
-        unsafe { &mut *self.world }
-            .resource_mut::<DeferredSystemRunQueue>()
-            .0
-            .push(self.id);
+        unsafe { &mut *self.run_queue }.push(self.id);
     }
 }
 
 unsafe impl Send for DeferredSystem {}
 unsafe impl Sync for DeferredSystem {}
 
-pub struct OnDropUnregisterDeferredSystem(pub DeferredSystem);
-
-impl Drop for OnDropUnregisterDeferredSystem {
-    fn drop(&mut self) {
-        // TODO: This is not sound. Pointer to world won't be valid across frames.
-        unsafe { &mut *self.0.world }
-            .remove_system(self.0.id)
-            .unwrap();
-    }
-}
-
-#[derive(Resource, Default)]
-pub struct DeferredSystemRunQueue(pub Vec<SystemId>);
+#[derive(Resource, Clone, Default)]
+pub struct DeferredSystemRunQueue(pub Box<Vec<SystemId>>);
diff --git a/src/hooks.rs b/src/hooks.rs
index 3db8862..77d69f9 100644
--- a/src/hooks.rs
+++ b/src/hooks.rs
@@ -1,7 +1,4 @@
-use crate::{
-    deferred_system::{DeferredSystem, OnDropUnregisterDeferredSystem},
-    tick::EcsContext,
-};
+use crate::{deferred_system::DeferredSystem, tick::EcsContext};
 use bevy::ecs::{
     system::{IntoSystem, Resource},
     world::World,
@@ -29,10 +26,7 @@ impl DioxusUiHooks for ScopeState {
     where
         S: IntoSystem<(), (), ()> + 'static,
     {
-        self.use_hook(|| {
-            OnDropUnregisterDeferredSystem(DeferredSystem::new(system, EcsContext::get_world(self)))
-        })
-        .0
+        *self.use_hook(|| DeferredSystem::new(system, EcsContext::get_world(self)))
     }
 }
 
diff --git a/src/tick.rs b/src/tick.rs
index 66616e8..a3f1cea 100644
--- a/src/tick.rs
+++ b/src/tick.rs
@@ -55,7 +55,7 @@ pub fn tick_dioxus_ui(world: &mut World) {
         );
     }
 
-    for system_id in mem::take(&mut world.resource_mut::<DeferredSystemRunQueue>().0) {
+    for system_id in mem::take(&mut *world.resource_mut::<DeferredSystemRunQueue>().0) {
         world.run_system(system_id).unwrap();
     }
 }