Cleanup deferred systems

src/deferred_system.rs

@@ -1,7 +1,11 @@
-use bevy::ecs::{
-    system::{IntoSystem, Resource, SystemId},
-    world::World,
+use bevy::{
+    ecs::{
+        system::{IntoSystem, Resource, SystemId},
+        world::World,
+    },
+    utils::HashMap,
 };
+use std::sync::Arc;
 
 #[derive(Clone, Copy)]
 pub struct DeferredSystem {
@@ -10,14 +14,23 @@ pub struct DeferredSystem {
 }
 
 impl DeferredSystem {
-    pub(crate) fn new<S>(system: S, world: &mut World) -> Self
+    pub(crate) fn new<S>(system: S, world: &mut World) -> (Self, Arc<()>)
     where
         S: IntoSystem<(), (), ()> + 'static,
     {
-        Self {
-            id: world.register_system(system), // TODO: We never unregister the system
-            run_queue: Box::as_mut(&mut world.resource_mut::<DeferredSystemRunQueue>().0),
-        }
+        let id = world.register_system(system);
+        let ref_count = Arc::new(());
+
+        let mut system_registry = world.resource_mut::<DeferredSystemRegistry>();
+        system_registry
+            .ref_counts
+            .insert(id, Arc::clone(&ref_count));
+
+        let deferred_system = Self {
+            id,
+            run_queue: Box::as_mut(&mut system_registry.run_queue),
+        };
+        (deferred_system, ref_count)
     }
 
     pub fn schedule(&self) {
@@ -28,5 +41,8 @@ impl DeferredSystem {
 unsafe impl Send for DeferredSystem {}
 unsafe impl Sync for DeferredSystem {}
 
-#[derive(Resource, Clone, Default)]
-pub struct DeferredSystemRunQueue(pub Box<Vec<SystemId>>);
+#[derive(Resource, Default)]
+pub struct DeferredSystemRegistry {
+    pub ref_counts: HashMap<SystemId, Arc<()>>,
+    pub run_queue: Box<Vec<SystemId>>,
+}

src/hooks.rs

@@ -26,7 +26,8 @@ impl DioxusUiHooks for ScopeState {
     where
         S: IntoSystem<(), (), ()> + 'static,
     {
-        *self.use_hook(|| DeferredSystem::new(system, EcsContext::get_world(self)))
+        self.use_hook(|| DeferredSystem::new(system, EcsContext::get_world(self)))
+            .0
     }
 }
 

src/lib.rs

@@ -6,7 +6,7 @@ mod tick;
 
 use self::{
     bsn::Bsn,
-    deferred_system::DeferredSystemRunQueue,
+    deferred_system::DeferredSystemRegistry,
     tick::{tick_dioxus_ui, VirtualDomUnsafe},
 };
 use bevy::{
@@ -23,7 +23,7 @@ pub struct DioxusUiPlugin;
 
 impl Plugin for DioxusUiPlugin {
     fn build(&self, app: &mut App) {
-        app.init_resource::<DeferredSystemRunQueue>()
+        app.init_resource::<DeferredSystemRegistry>()
             .add_systems(Update, tick_dioxus_ui);
     }
 }

src/tick.rs

@@ -1,13 +1,16 @@
 use crate::{
-    apply_mutations::apply_mutations, deferred_system::DeferredSystemRunQueue, DioxusUiRoot,
+    apply_mutations::apply_mutations, deferred_system::DeferredSystemRegistry, DioxusUiRoot,
 };
 use bevy::{
-    ecs::{entity::Entity, world::World},
+    ecs::{
+        entity::Entity,
+        world::{Mut, World},
+    },
     prelude::{Deref, DerefMut},
     utils::synccell::SyncCell,
 };
 use dioxus_core::{Element, Scope, ScopeState, VirtualDom};
-use std::mem;
+use std::{mem, sync::Arc};
 
 // TODO: This is not sound. Can't borrow the world while iterating over DioxusUiRoots.
 pub fn tick_dioxus_ui(world: &mut World) {
@@ -55,9 +58,19 @@ pub fn tick_dioxus_ui(world: &mut World) {
         );
     }
 
-    for system_id in mem::take(&mut *world.resource_mut::<DeferredSystemRunQueue>().0) {
-        world.run_system(system_id).unwrap();
+    for system_id in mem::take(&mut *world.resource_mut::<DeferredSystemRegistry>().run_queue) {
+        let _ = world.run_system(system_id);
     }
+
+    world.resource_scope(|world, mut system_registry: Mut<DeferredSystemRegistry>| {
+        system_registry.ref_counts.retain(|system_id, ref_count| {
+            let cleanup = Arc::strong_count(ref_count) == 1;
+            if cleanup {
+                world.remove_system(*system_id).unwrap();
+            }
+            !cleanup
+        });
+    });
 }
 
 #[derive(Clone)]