Move VirtualDom to a non-send resource

src/lib.rs

@@ -6,19 +6,17 @@ pub mod hooks;
 mod tick;
 
 use self::{
-    apply_mutations::BevyTemplate,
+    apply_mutations::BevyTemplate, deferred_system::DeferredSystemRegistry, events::EventReaders,
-    deferred_system::DeferredSystemRegistry,
+    hooks::EcsSubscriptions, tick::tick_dioxus_ui,
-    events::EventReaders,
-    hooks::EcsSubscriptions,
-    tick::{tick_dioxus_ui, VirtualDomUnsafe},
 };
 use bevy::{
     app::{App, Plugin, Update},
     ecs::{bundle::Bundle, component::Component, entity::Entity},
+    prelude::{Deref, DerefMut},
     ui::node_bundles::NodeBundle,
     utils::{EntityHashMap, HashMap},
 };
-use dioxus::core::{Element, ElementId, Scope};
+use dioxus::core::{Element, ElementId, Scope, VirtualDom};
 
 pub use bevy_mod_picking;
 pub use dioxus;
@@ -27,7 +25,8 @@ pub struct DioxusUiPlugin;
 
 impl Plugin for DioxusUiPlugin {
     fn build(&self, app: &mut App) {
-        app.init_resource::<EcsSubscriptions>()
+        app.init_non_send_resource::<UiRoots>()
+            .init_resource::<EcsSubscriptions>()
             .init_resource::<DeferredSystemRegistry>()
             .init_resource::<EventReaders>()
             .add_systems(Update, tick_dioxus_ui);
@@ -40,19 +39,30 @@ pub struct DioxusUiBundle {
     pub node_bundle: NodeBundle,
 }
 
-#[derive(Component)]
+#[derive(Component, Deref, Clone, Copy)]
-pub struct DioxusUiRoot {
+pub struct DioxusUiRoot(fn(Scope) -> Element);
-    virtual_dom: VirtualDomUnsafe,
+
+impl DioxusUiRoot {
+    pub fn new(root_component: fn(Scope) -> Element) -> Self {
+        Self(root_component)
+    }
+}
+
+#[derive(Deref, DerefMut, Default)]
+struct UiRoots(EntityHashMap<Entity, UiRoot>);
+
+struct UiRoot {
+    virtual_dom: VirtualDom,
     element_id_to_bevy_ui_entity: HashMap<ElementId, Entity>,
     bevy_ui_entity_to_element_id: EntityHashMap<Entity, ElementId>,
     templates: HashMap<String, BevyTemplate>,
     needs_rebuild: bool,
 }
 
-impl DioxusUiRoot {
+impl UiRoot {
-    pub fn new(root_component: fn(Scope) -> Element) -> Self {
+    fn new(root_component: fn(Scope) -> Element) -> Self {
         Self {
-            virtual_dom: VirtualDomUnsafe::new(root_component),
+            virtual_dom: VirtualDom::new(root_component),
             element_id_to_bevy_ui_entity: HashMap::new(),
             bevy_ui_entity_to_element_id: EntityHashMap::default(),
             templates: HashMap::new(),

src/tick.rs

@@ -3,19 +3,16 @@ use crate::{
     deferred_system::DeferredSystemRegistry,
     events::EventReaders,
     hooks::{EcsContext, EcsSubscriptions},
-    DioxusUiRoot,
+    DioxusUiRoot, UiRoot, UiRoots,
 };
 use bevy::{
     ecs::{
         entity::Entity,
-        query::With,
         world::{Mut, World},
     },
     hierarchy::Parent,
-    prelude::{Deref, DerefMut},
+    utils::HashMap,
-    utils::synccell::SyncCell,
 };
-use dioxus::core::{Element, Scope, VirtualDom};
 use std::{any::Any, mem, rc::Rc, sync::Arc};
 
 pub fn tick_dioxus_ui(world: &mut World) {
@@ -24,17 +21,32 @@ pub fn tick_dioxus_ui(world: &mut World) {
     let ui_events = world.resource_scope(|world, mut event_readers: Mut<EventReaders>| {
         event_readers.get_dioxus_events(world.resource())
     });
-    let root_entities: Vec<Entity> = world
+
-        .query_filtered::<Entity, With<DioxusUiRoot>>()
+    let root_entities: HashMap<Entity, DioxusUiRoot> = world
+        .query::<(Entity, &DioxusUiRoot)>()
         .iter(world)
+        .map(|(entity, ui_root)| (entity, *ui_root))
         .collect();
 
-    for root_entity in root_entities {
+    world
-        dispatch_ui_events(&ui_events, root_entity, world);
+        .non_send_resource_mut::<UiRoots>()
+        .retain(|root_entity, _| root_entities.contains_key(root_entity));
 
-        schedule_ui_renders_from_ecs_subscriptions(root_entity, world);
+    for (root_entity, ui_root) in root_entities {
+        let mut ui_root = world
+            .non_send_resource_mut::<UiRoots>()
+            .remove(&root_entity)
+            .unwrap_or_else(|| UiRoot::new(*ui_root));
 
-        render_ui(root_entity, world);
+        dispatch_ui_events(&ui_events, &mut ui_root, world);
+
+        schedule_ui_renders_from_ecs_subscriptions(&mut ui_root, world);
+
+        render_ui(root_entity, &mut ui_root, world);
+
+        world
+            .non_send_resource_mut::<UiRoots>()
+            .insert(root_entity, ui_root);
     }
 }
 
@@ -56,14 +68,9 @@ fn run_deferred_systems(world: &mut World) {
 
 fn dispatch_ui_events(
     events: &Vec<(Entity, &str, Rc<dyn Any>)>,
-    root_entity: Entity,
+    ui_root: &mut UiRoot,
-    world: &mut World,
+    world: &World,
 ) {
-    let mut ui_root = world
-        .entity_mut(root_entity)
-        .take::<DioxusUiRoot>()
-        .unwrap();
-
     for (mut target, name, data) in events {
         let mut target_element_id = ui_root.bevy_ui_entity_to_element_id.get(&target).copied();
         while target_element_id.is_none() {
@@ -71,26 +78,14 @@ fn dispatch_ui_events(
             target_element_id = ui_root.bevy_ui_entity_to_element_id.get(&target).copied();
         }
 
-        ui_root.virtual_dom.get().handle_event(
+        ui_root
-            name,
-            Rc::clone(data),
-            target_element_id.unwrap(),
-            true,
-        );
-    }
-
-    world.entity_mut(root_entity).insert(ui_root);
-}
-
-fn schedule_ui_renders_from_ecs_subscriptions(root_entity: Entity, world: &mut World) {
-    let schedule_update = world
-        .get_mut::<DioxusUiRoot>(root_entity)
-        .unwrap()
             .virtual_dom
-        .get()
+            .handle_event(name, Rc::clone(data), target_element_id.unwrap(), true);
-        .base_scope()
+    }
-        .schedule_update_any();
+}
 
+fn schedule_ui_renders_from_ecs_subscriptions(ui_root: &mut UiRoot, world: &World) {
+    let schedule_update = ui_root.virtual_dom.base_scope().schedule_update_any();
     let ecs_subscriptions = world.resource::<EcsSubscriptions>();
 
     for scope_id in &*ecs_subscriptions.world_and_queries {
@@ -106,21 +101,15 @@ fn schedule_ui_renders_from_ecs_subscriptions(root_entity: Entity, world: &mut W
     }
 }
 
-fn render_ui(root_entity: Entity, world: &mut World) {
+fn render_ui(root_entity: Entity, ui_root: &mut UiRoot, world: &mut World) {
-    let mut ui_root = world
-        .entity_mut(root_entity)
-        .take::<DioxusUiRoot>()
-        .unwrap();
-
     ui_root
         .virtual_dom
-        .get()
         .base_scope()
         .provide_context(EcsContext { world });
 
     if ui_root.needs_rebuild {
         apply_mutations(
-            ui_root.virtual_dom.get().rebuild(),
+            ui_root.virtual_dom.rebuild(),
             &mut ui_root.element_id_to_bevy_ui_entity,
             &mut ui_root.bevy_ui_entity_to_element_id,
             &mut ui_root.templates,
@@ -131,23 +120,11 @@ fn render_ui(root_entity: Entity, world: &mut World) {
     }
 
     apply_mutations(
-        ui_root.virtual_dom.get().render_immediate(),
+        ui_root.virtual_dom.render_immediate(),
         &mut ui_root.element_id_to_bevy_ui_entity,
         &mut ui_root.bevy_ui_entity_to_element_id,
         &mut ui_root.templates,
         root_entity,
         world,
     );
-
-    world.entity_mut(root_entity).insert(ui_root);
-}
-
-#[derive(Deref, DerefMut)]
-pub struct VirtualDomUnsafe(pub SyncCell<VirtualDom>);
-unsafe impl Send for VirtualDomUnsafe {}
-
-impl VirtualDomUnsafe {
-    pub fn new(root_component: fn(Scope) -> Element) -> Self {
-        Self(SyncCell::new(VirtualDom::new(root_component)))
-    }
 }