diff --git a/src/deferred_system.rs b/src/deferred_system.rs index 16170ac..85da115 100644 --- a/src/deferred_system.rs +++ b/src/deferred_system.rs @@ -2,16 +2,10 @@ use bevy::ecs::{ system::{IntoSystem, Resource, SystemId}, world::World, }; -use std::sync::Arc; -// TODO: Can probably make the value stored in the hook the only one that has a Drop unregistering the system, -// and then make the type that lets you schedule the system Copy - -#[derive(Clone)] -pub struct DeferredSystem(pub(crate) Arc); - -pub(crate) struct DeferredSystemInner { - pub id: SystemId, +#[derive(Clone, Copy)] +pub struct DeferredSystem { + id: SystemId, world: *mut World, } @@ -20,28 +14,32 @@ impl DeferredSystem { where S: IntoSystem<(), (), ()> + 'static, { - Self(Arc::new(DeferredSystemInner { + Self { id: world.register_system(system), world, - })) + } } pub fn schedule(&self) { - unsafe { &mut *self.0.world } + unsafe { &mut *self.world } .resource_mut::() .0 - .push(self.clone()); + .push(self.id); } } -impl Drop for DeferredSystemInner { +unsafe impl Send for DeferredSystem {} +unsafe impl Sync for DeferredSystem {} + +pub struct OnDropUnregisterDeferredSystem(pub DeferredSystem); + +impl Drop for OnDropUnregisterDeferredSystem { fn drop(&mut self) { - unsafe { &mut *self.world }.remove_system(self.id).unwrap(); + unsafe { &mut *self.0.world } + .remove_system(self.0.id) + .unwrap(); } } -unsafe impl Send for DeferredSystemInner {} -unsafe impl Sync for DeferredSystemInner {} - #[derive(Resource, Default)] -pub struct DeferredSystemRunQueue(pub Vec); +pub struct DeferredSystemRunQueue(pub Vec); diff --git a/src/hooks.rs b/src/hooks.rs index 6ee8901..3db8862 100644 --- a/src/hooks.rs +++ b/src/hooks.rs @@ -1,4 +1,7 @@ -use crate::{deferred_system::DeferredSystem, tick::EcsContext}; +use crate::{ + deferred_system::{DeferredSystem, OnDropUnregisterDeferredSystem}, + tick::EcsContext, +}; use bevy::ecs::{ system::{IntoSystem, Resource}, world::World, @@ -26,8 +29,10 @@ impl DioxusUiHooks for ScopeState { where S: IntoSystem<(), (), ()> + 'static, { - self.use_hook(|| DeferredSystem::new(system, EcsContext::get_world(self))) - .clone() + self.use_hook(|| { + OnDropUnregisterDeferredSystem(DeferredSystem::new(system, EcsContext::get_world(self))) + }) + .0 } } diff --git a/src/tick.rs b/src/tick.rs index 3e24d7a..0d0369f 100644 --- a/src/tick.rs +++ b/src/tick.rs @@ -44,8 +44,8 @@ pub fn tick_dioxus_ui(world: &mut World) { }); } - for system in mem::take(&mut world.resource_mut::().0) { - world.run_system(system.0.id).unwrap(); + for system_id in mem::take(&mut world.resource_mut::().0) { + world.run_system(system_id).unwrap(); } }