From 99929cbc7fa20d4d95a8f4e12dbdf97b6489f603 Mon Sep 17 00:00:00 2001
From: JMS55 <47158642+JMS55@users.noreply.github.com>
Date: Tue, 12 Dec 2023 16:22:54 -0800
Subject: [PATCH] Make ECS hooks subscribe to changes

---
 src/hooks.rs | 71 ++++++++++++++++++++++++++++++++++++++++++++++------
 src/lib.rs   |  4 ++-
 src/tick.rs  | 16 +++++++++++-
 3 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/src/hooks.rs b/src/hooks.rs
index a0e896a..2f7e38c 100644
--- a/src/hooks.rs
+++ b/src/hooks.rs
@@ -2,21 +2,65 @@ use crate::{
     deferred_system::{new_deferred_system, DeferredSystem},
     tick::EcsContext,
 };
-use bevy::ecs::{
-    query::{QueryState, ReadOnlyWorldQuery},
-    system::{IntoSystem, Query, Resource},
-    world::{unsafe_world_cell::UnsafeWorldCell, World},
+use bevy::{
+    ecs::{
+        component::ComponentId,
+        query::{QueryState, ReadOnlyWorldQuery},
+        system::{IntoSystem, Query, Resource},
+        world::{unsafe_world_cell::UnsafeWorldCell, World},
+    },
+    utils::{HashMap, HashSet},
+};
+use dioxus::{
+    core::{ScopeId, ScopeState},
+    hooks::use_on_destroy,
 };
-use dioxus::core::ScopeState;
 
-// TODO: Hooks need to schedule future updates
+#[derive(Resource, Default)]
+pub(crate) struct EcsSubscriptions {
+    pub resources: Box<HashMap<ComponentId, HashSet<ScopeId>>>,
+    pub world_and_queries: Box<HashSet<ScopeId>>,
+}
 
 pub fn use_world<'a>(cx: &'a ScopeState) -> &'a World {
-    EcsContext::get_world(cx)
+    let world = EcsContext::get_world(cx);
+
+    let scope_id = cx.scope_id();
+    let subscription_manager = *cx.use_hook(|| {
+        let subscription_manager = &mut world.resource_mut::<EcsSubscriptions>().world_and_queries;
+        subscription_manager.insert(scope_id);
+        Box::as_mut(subscription_manager) as *mut HashSet<ScopeId>
+    });
+    use_on_destroy(cx, move || {
+        unsafe { &mut *subscription_manager }.remove(&scope_id);
+    });
+
+    world
 }
 
 pub fn use_resource<'a, T: Resource>(cx: &'a ScopeState) -> &'a T {
-    EcsContext::get_world(cx).resource()
+    let world = EcsContext::get_world(cx);
+
+    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;
+        subscription_manager
+            .entry(resource_id)
+            .or_default()
+            .insert(scope_id);
+        Box::as_mut(subscription_manager) as *mut HashMap<ComponentId, HashSet<ScopeId>>
+    });
+    use_on_destroy(cx, move || {
+        let subscription_manager = &mut unsafe { &mut *subscription_manager };
+        let resource_subscriptions = subscription_manager.get_mut(&resource_id).unwrap();
+        resource_subscriptions.remove(&scope_id);
+        if resource_subscriptions.is_empty() {
+            subscription_manager.remove(&resource_id);
+        }
+    });
+
+    world.resource()
 }
 
 pub fn use_query<'a, Q>(cx: &'a ScopeState) -> DioxusUiQuery<'a, Q, ()>
@@ -32,6 +76,17 @@ where
     F: ReadOnlyWorldQuery,
 {
     let world = EcsContext::get_world(cx);
+
+    let scope_id = cx.scope_id();
+    let subscription_manager = *cx.use_hook(|| {
+        let subscription_manager = &mut world.resource_mut::<EcsSubscriptions>().world_and_queries;
+        subscription_manager.insert(scope_id);
+        Box::as_mut(subscription_manager) as *mut HashSet<ScopeId>
+    });
+    use_on_destroy(cx, move || {
+        unsafe { &mut *subscription_manager }.remove(&scope_id);
+    });
+
     DioxusUiQuery {
         query_state: QueryState::new(world),
         world_cell: world.as_unsafe_world_cell(),
diff --git a/src/lib.rs b/src/lib.rs
index caa075c..55ca3fe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,6 +8,7 @@ use self::{
     apply_mutations::BevyTemplate,
     deferred_system::DeferredSystemRegistry,
     events::EventReaders,
+    hooks::EcsSubscriptions,
     tick::{tick_dioxus_ui, VirtualDomUnsafe},
 };
 use bevy::{
@@ -26,7 +27,8 @@ pub struct DioxusUiPlugin;
 
 impl Plugin for DioxusUiPlugin {
     fn build(&self, app: &mut App) {
-        app.init_resource::<DeferredSystemRegistry>()
+        app.init_resource::<EcsSubscriptions>()
+            .init_resource::<DeferredSystemRegistry>()
             .init_resource::<EventReaders>()
             .add_systems(Update, tick_dioxus_ui);
     }
diff --git a/src/tick.rs b/src/tick.rs
index ed8a581..e9354ce 100644
--- a/src/tick.rs
+++ b/src/tick.rs
@@ -1,6 +1,6 @@
 use crate::{
     apply_mutations::apply_mutations, deferred_system::DeferredSystemRegistry,
-    events::EventReaders, DioxusUiRoot,
+    events::EventReaders, hooks::EcsSubscriptions, DioxusUiRoot,
 };
 use bevy::{
     ecs::{
@@ -17,6 +17,7 @@ use std::{mem, rc::Rc, sync::Arc};
 pub fn tick_dioxus_ui(world: &mut World) {
     let world_ptr: *mut World = world;
     let world_cell = world.as_unsafe_world_cell();
+    let ecs_subscriptions = unsafe { world_cell.get_resource::<EcsSubscriptions>().unwrap() };
     let events = unsafe {
         world_cell
             .get_resource_mut::<EventReaders>()
@@ -42,6 +43,19 @@ pub fn tick_dioxus_ui(world: &mut World) {
         } = &mut *dioxus_ui_root;
         let virtual_dom = virtual_dom.get();
 
+        let schedule_update = virtual_dom.base_scope().schedule_update_any();
+        for scope_id in &*ecs_subscriptions.world_and_queries {
+            schedule_update(*scope_id);
+        }
+        for (resource_id, scope_ids) in &*ecs_subscriptions.resources {
+            // TODO: Only schedule updates if resource has changed
+            // if unsafe { world_cell.world() }.is_resource_changed() {
+            for scope_id in scope_ids {
+                schedule_update(*scope_id);
+            }
+            // }
+        }
+
         virtual_dom
             .base_scope()
             .provide_context(EcsContext { world: world_ptr });