From fa03c65f8ae26979b4d7e6a12a16de4a7fd1593d Mon Sep 17 00:00:00 2001
From: JMS55 <47158642+JMS55@users.noreply.github.com>
Date: Sat, 16 Dec 2023 20:19:41 -0800
Subject: [PATCH] Fix segfault

---
 src/hooks.rs | 19 ++++++++++++++-----
 src/lib.rs   | 13 +++++++------
 src/tick.rs  | 14 ++++++--------
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/src/hooks.rs b/src/hooks.rs
index 4443d38..161dbba 100644
--- a/src/hooks.rs
+++ b/src/hooks.rs
@@ -1,4 +1,4 @@
-use crate::deferred_system::new_deferred_system;
+use crate::{deferred_system::new_deferred_system, UiContext};
 use bevy::{
     ecs::{
         component::ComponentId,
@@ -13,7 +13,7 @@ use dioxus::{
     hooks::use_on_destroy,
 };
 
-#[derive(Resource, Default)]
+#[derive(Default)]
 pub(crate) struct EcsSubscriptions {
     pub resources: Box<HashMap<ComponentId, HashSet<ScopeId>>>,
     pub world_and_queries: Box<HashSet<ScopeId>>,
@@ -40,7 +40,10 @@ pub fn use_world<'a>(cx: &'a ScopeState) -> &'a World {
 
     let scope_id = cx.scope_id();
     let subscription_manager = *cx.use_hook(|| {
-        let subscription_manager = &mut world.resource_mut::<EcsSubscriptions>().world_and_queries;
+        let subscription_manager = &mut world
+            .non_send_resource_mut::<UiContext>()
+            .subscriptions
+            .world_and_queries;
         subscription_manager.insert(scope_id);
         Box::as_mut(subscription_manager) as *mut HashSet<ScopeId>
     });
@@ -57,7 +60,10 @@ pub fn use_resource<'a, T: Resource>(cx: &'a ScopeState) -> &'a T {
     let resource_id = world.components().resource_id::<T>().unwrap();
     let scope_id = cx.scope_id();
     let subscription_manager = *cx.use_hook(|| {
-        let subscription_manager = &mut world.resource_mut::<EcsSubscriptions>().resources;
+        let subscription_manager = &mut world
+            .non_send_resource_mut::<UiContext>()
+            .subscriptions
+            .resources;
         subscription_manager
             .entry(resource_id)
             .or_default()
@@ -92,7 +98,10 @@ where
 
     let scope_id = cx.scope_id();
     let subscription_manager = *cx.use_hook(|| {
-        let subscription_manager = &mut world.resource_mut::<EcsSubscriptions>().world_and_queries;
+        let subscription_manager = &mut world
+            .non_send_resource_mut::<UiContext>()
+            .subscriptions
+            .world_and_queries;
         subscription_manager.insert(scope_id);
         Box::as_mut(subscription_manager) as *mut HashSet<ScopeId>
     });
diff --git a/src/lib.rs b/src/lib.rs
index 6982d65..6635d34 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -12,7 +12,7 @@ use self::{
 use bevy::{
     app::{App, Plugin, Update},
     ecs::{bundle::Bundle, component::Component, entity::Entity},
-    prelude::{Deref, DerefMut},
+    prelude::Deref,
     ui::node_bundles::NodeBundle,
     utils::{EntityHashMap, HashMap},
 };
@@ -25,9 +25,7 @@ pub struct DioxusUiPlugin;
 
 impl Plugin for DioxusUiPlugin {
     fn build(&self, app: &mut App) {
-        // TODO: I think UiRoots must be dropped only after EcsSubscriptions
-        app.init_non_send_resource::<UiRoots>()
-            .init_resource::<EcsSubscriptions>()
+        app.init_non_send_resource::<UiContext>()
             .init_resource::<DeferredSystemRegistry>()
             .init_resource::<EventReaders>()
             .add_systems(Update, tick_dioxus_ui);
@@ -43,8 +41,11 @@ pub struct DioxusUiBundle {
 #[derive(Component, Deref, Hash, PartialEq, Eq, Clone, Copy)]
 pub struct DioxusUiRoot(pub fn(Scope) -> Element);
 
-#[derive(Deref, DerefMut, Default)]
-struct UiRoots(HashMap<(Entity, DioxusUiRoot), UiRoot>);
+#[derive(Default)]
+struct UiContext {
+    roots: HashMap<(Entity, DioxusUiRoot), UiRoot>,
+    subscriptions: EcsSubscriptions,
+}
 
 struct UiRoot {
     virtual_dom: VirtualDom,
diff --git a/src/tick.rs b/src/tick.rs
index 1b5081b..5308e4e 100644
--- a/src/tick.rs
+++ b/src/tick.rs
@@ -1,9 +1,6 @@
 use crate::{
-    apply_mutations::apply_mutations,
-    deferred_system::DeferredSystemRegistry,
-    events::EventReaders,
-    hooks::{EcsContext, EcsSubscriptions},
-    DioxusUiRoot, UiRoot, UiRoots,
+    apply_mutations::apply_mutations, deferred_system::DeferredSystemRegistry,
+    events::EventReaders, hooks::EcsContext, DioxusUiRoot, UiContext, UiRoot,
 };
 use bevy::{
     ecs::{
@@ -28,7 +25,7 @@ pub fn tick_dioxus_ui(world: &mut World) {
         .iter(world)
         .map(|(entity, ui_root)| (entity, *ui_root))
         .collect();
-    let mut ui_roots = mem::take(&mut world.non_send_resource_mut::<UiRoots>().0);
+    let mut ui_roots = mem::take(&mut world.non_send_resource_mut::<UiContext>().roots);
 
     for (root_entity, dioxus_ui_root) in root_entities {
         let mut ui_root = ui_roots
@@ -42,7 +39,8 @@ pub fn tick_dioxus_ui(world: &mut World) {
         render_ui(root_entity, &mut ui_root, world);
 
         world
-            .non_send_resource_mut::<UiRoots>()
+            .non_send_resource_mut::<UiContext>()
+            .roots
             .insert((root_entity, dioxus_ui_root), ui_root);
     }
 }
@@ -82,7 +80,7 @@ fn dispatch_ui_events(
 }
 
 fn schedule_ui_renders_from_ecs_subscriptions(ui_root: &mut UiRoot, world: &World) {
-    let ecs_subscriptions = world.resource::<EcsSubscriptions>();
+    let ecs_subscriptions = &world.non_send_resource::<UiContext>().subscriptions;
 
     for scope_id in &*ecs_subscriptions.world_and_queries {
         ui_root.virtual_dom.mark_dirty(*scope_id);