diff --git a/examples/demo.rs b/examples/demo.rs
index 7ab1d03..87878de 100644
--- a/examples/demo.rs
+++ b/examples/demo.rs
@@ -16,7 +16,7 @@ fn main() {
         .add_plugins((DefaultPlugins, DioxusUiPlugin, DefaultPickingPlugins))
         .add_systems(Startup, |mut commands: Commands| {
             commands.spawn(DioxusUiBundle {
-                dioxus_ui_root: DioxusUiRoot::new(Editor),
+                dioxus_ui_root: DioxusUiRoot(Editor),
                 node_bundle: NodeBundle::default(),
             });
             commands.spawn(Camera2dBundle::default());
diff --git a/src/lib.rs b/src/lib.rs
index 2aa2ac1..845c6c5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -39,17 +39,11 @@ pub struct DioxusUiBundle {
     pub node_bundle: NodeBundle,
 }
 
-#[derive(Component, Deref, Clone, Copy)]
-pub struct DioxusUiRoot(fn(Scope) -> Element);
-
-impl DioxusUiRoot {
-    pub fn new(root_component: fn(Scope) -> Element) -> Self {
-        Self(root_component)
-    }
-}
+#[derive(Component, Deref, Hash, PartialEq, Eq, Clone, Copy)]
+pub struct DioxusUiRoot(pub fn(Scope) -> Element);
 
 #[derive(Deref, DerefMut, Default)]
-struct UiRoots(EntityHashMap<Entity, UiRoot>);
+struct UiRoots(HashMap<(Entity, DioxusUiRoot), UiRoot>);
 
 struct UiRoot {
     virtual_dom: VirtualDom,
@@ -60,9 +54,9 @@ struct UiRoot {
 }
 
 impl UiRoot {
-    fn new(root_component: fn(Scope) -> Element) -> Self {
+    fn new(root_component: DioxusUiRoot) -> Self {
         Self {
-            virtual_dom: VirtualDom::new(root_component),
+            virtual_dom: VirtualDom::new(root_component.0),
             element_id_to_bevy_ui_entity: HashMap::new(),
             bevy_ui_entity_to_element_id: EntityHashMap::default(),
             templates: HashMap::new(),
diff --git a/src/tick.rs b/src/tick.rs
index 70a7b5c..1352251 100644
--- a/src/tick.rs
+++ b/src/tick.rs
@@ -27,16 +27,12 @@ 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);
 
-    world
-        .non_send_resource_mut::<UiRoots>()
-        .retain(|root_entity, _| root_entities.contains_key(root_entity));
-
-    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));
+    for (root_entity, dioxus_ui_root) in root_entities {
+        let mut ui_root = ui_roots
+            .remove(&(root_entity, dioxus_ui_root))
+            .unwrap_or_else(|| UiRoot::new(dioxus_ui_root));
 
         dispatch_ui_events(&ui_events, &mut ui_root, world);
 
@@ -46,7 +42,7 @@ pub fn tick_dioxus_ui(world: &mut World) {
 
         world
             .non_send_resource_mut::<UiRoots>()
-            .insert(root_entity, ui_root);
+            .insert((root_entity, dioxus_ui_root), ui_root);
     }
 }