rerun tasks in the same order as components

packages/core/src/arena.rs

@@ -1,4 +1,5 @@
-use crate::{innerlude::DirtyScope, virtual_dom::VirtualDom, ScopeId};
+use crate::innerlude::ScopeOrder;
+use crate::{virtual_dom::VirtualDom, ScopeId};
 
 /// An Element's unique identifier.
 ///
@@ -74,7 +75,7 @@ impl VirtualDom {
             context.height
         };
 
-        self.dirty_scopes.remove(&DirtyScope { height, id });
+        self.dirty_scopes.remove(&ScopeOrder::new(height, id));
     }
 }
 

packages/core/src/diff/component.rs

@@ -91,10 +91,7 @@ impl VNode {
         dom.diff_scope(to, scope_id, new);
 
         let height = dom.runtime.get_state(scope_id).unwrap().height;
-        dom.dirty_scopes.remove(&DirtyScope {
-            height,
-            id: scope_id,
-        });
+        dom.dirty_scopes.remove(&DirtyScope::new(height, scope_id));
     }
 
     fn replace_vcomponent(

packages/core/src/dirty_scope.rs

@@ -1,33 +1,103 @@
+use crate::ScopeId;
+use crate::Task;
+use std::borrow::Borrow;
+use std::cell::Cell;
+use std::cell::RefCell;
 use std::hash::Hash;
 
-use crate::ScopeId;
-
 #[derive(Debug, Clone, Eq)]
-pub struct DirtyScope {
-    pub height: u32,
-    pub id: ScopeId,
+pub struct ScopeOrder {
+    pub(crate) height: u32,
+    pub(crate) id: ScopeId,
 }
 
-impl PartialOrd for DirtyScope {
-    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-        Some(self.cmp(other))
+impl ScopeOrder {
+    pub fn new(height: u32, id: ScopeId) -> Self {
+        Self { height, id }
     }
 }
 
-impl Ord for DirtyScope {
-    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-        self.height.cmp(&other.height).then(self.id.cmp(&other.id))
-    }
-}
-
-impl PartialEq for DirtyScope {
+impl PartialEq for ScopeOrder {
     fn eq(&self, other: &Self) -> bool {
         self.id == other.id
     }
 }
 
-impl Hash for DirtyScope {
+impl PartialOrd for ScopeOrder {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for ScopeOrder {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.height.cmp(&other.height).then(self.id.cmp(&other.id))
+    }
+}
+
+impl Hash for ScopeOrder {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
         self.id.hash(state);
     }
 }
+
+#[derive(Debug, Clone, Eq)]
+pub struct DirtyScope {
+    pub order: ScopeOrder,
+    pub rerun_queued: Cell<bool>,
+    pub tasks_queued: RefCell<Vec<Task>>,
+}
+
+impl From<ScopeOrder> for DirtyScope {
+    fn from(order: ScopeOrder) -> Self {
+        Self {
+            order,
+            rerun_queued: false.into(),
+            tasks_queued: Vec::new().into(),
+        }
+    }
+}
+
+impl DirtyScope {
+    pub fn new(height: u32, id: ScopeId) -> Self {
+        ScopeOrder { height, id }.into()
+    }
+
+    pub fn queue_task(&self, task: Task) {
+        self.tasks_queued.borrow_mut().push(task);
+    }
+
+    pub fn queue_rerun(&self) {
+        self.rerun_queued.set(true);
+    }
+}
+
+impl Borrow<ScopeOrder> for DirtyScope {
+    fn borrow(&self) -> &ScopeOrder {
+        &self.order
+    }
+}
+
+impl PartialOrd for DirtyScope {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.order.cmp(&other.order))
+    }
+}
+
+impl Ord for DirtyScope {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.order.cmp(&other.order)
+    }
+}
+
+impl PartialEq for DirtyScope {
+    fn eq(&self, other: &Self) -> bool {
+        self.order == other.order
+    }
+}
+
+impl Hash for DirtyScope {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.order.hash(state);
+    }
+}

packages/core/src/scope_arena.rs

@@ -1,6 +1,7 @@
+use crate::innerlude::ScopeOrder;
 use crate::{
     any_props::{AnyProps, BoxedAnyProps},
-    innerlude::{DirtyScope, ScopeState},
+    innerlude::ScopeState,
     nodes::RenderReturn,
     scope_context::Scope,
     scopes::ScopeId,
@@ -66,10 +67,8 @@ impl VirtualDom {
         context.render_count.set(context.render_count.get() + 1);
 
         // remove this scope from dirty scopes
-        self.dirty_scopes.remove(&DirtyScope {
-            height: context.height,
-            id: context.id,
-        });
+        self.dirty_scopes
+            .remove(&ScopeOrder::new(context.height, scope_id));
 
         if context.suspended.get() {
             if matches!(new_nodes, RenderReturn::Aborted(_)) {

packages/core/src/tasks.rs

@@ -135,6 +135,10 @@ impl Runtime {
         self.tasks.borrow().get(task.0)?.parent
     }
 
+    pub(crate) fn task_scope(&self, task: Task) -> Option<ScopeId> {
+        self.tasks.borrow().get(task.0).map(|t| t.scope)
+    }
+
     pub(crate) fn handle_task_wakeup(&self, id: Task) -> Poll<()> {
         debug_assert!(Runtime::current().is_some(), "Must be in a dioxus runtime");
 

packages/core/src/virtual_dom.rs

@@ -2,6 +2,8 @@
 //!
 //! This module provides the primary mechanics to create a hook-based, concurrent VDOM for Rust.
 
+use crate::innerlude::ScopeOrder;
+use crate::Task;
 use crate::{
     any_props::AnyProps,
     arena::ElementId,
@@ -183,6 +185,7 @@ pub struct VirtualDom {
     pub(crate) scopes: Slab<ScopeState>,
 
     pub(crate) dirty_scopes: BTreeSet<DirtyScope>,
+    pub(crate) scopes_need_rerun: bool,
 
     // Maps a template path to a map of byte indexes to templates
     pub(crate) templates: FxHashMap<TemplateId, FxHashMap<usize, Template>>,
@@ -310,6 +313,7 @@ impl VirtualDom {
             rx,
             runtime: Runtime::new(tx),
             scopes: Default::default(),
+            scopes_need_rerun: false,
             dirty_scopes: Default::default(),
             templates: Default::default(),
             queued_templates: Default::default(),
@@ -374,10 +378,39 @@ impl VirtualDom {
         };
 
         tracing::trace!("Marking scope {:?} as dirty", id);
-        self.dirty_scopes.insert(DirtyScope {
-            height: scope.height(),
-            id,
-        });
+        self.scopes_need_rerun = true;
+        let order = ScopeOrder::new(scope.height(), id);
+        match self.dirty_scopes.get(&order) {
+            Some(dirty) => {
+                dirty.queue_rerun();
+            }
+            None => {
+                let dirty: DirtyScope = order.into();
+                dirty.queue_rerun();
+                self.dirty_scopes.insert(dirty);
+            }
+        }
+    }
+
+    /// Mark a task as dirty
+    fn mark_task_dirty(&mut self, task: Task) {
+        let Some(scope) = self.runtime.task_scope(task) else {
+            return;
+        };
+        let Some(scope) = self.runtime.get_state(scope) else {
+            return;
+        };
+        let order = ScopeOrder::new(scope.height(), scope.id);
+        match self.dirty_scopes.get(&order) {
+            Some(dirty) => {
+                dirty.queue_task(task);
+            }
+            None => {
+                let dirty: DirtyScope = order.into();
+                dirty.queue_task(task);
+                self.dirty_scopes.insert(dirty);
+            }
+        }
     }
 
     /// Call a listener inside the VirtualDom with data from outside the VirtualDom. **The ElementId passed in must be the id of an element with a listener, not a static node or a text node.**
@@ -439,13 +472,32 @@ impl VirtualDom {
             self.process_events();
 
             // Now that we have collected all queued work, we should check if we have any dirty scopes. If there are not, then we can poll any queued futures
-            if !self.dirty_scopes.is_empty() {
+            if self.scopes_need_rerun {
                 return;
             }
 
             // Make sure we set the runtime since we're running user code
             let _runtime = RuntimeGuard::new(self.runtime.clone());
 
+            // Next, run any queued tasks
+            // We choose not to poll the deadline since we complete pretty quickly anyways
+            while let Some(dirty) = self.dirty_scopes.pop_first() {
+                // If the scope doesn't exist for whatever reason, then we should skip it
+                if !self.scopes.contains(dirty.order.id.0) {
+                    continue;
+                }
+
+                // Then poll any tasks that might be pending
+                for task in dirty.tasks_queued.borrow().iter() {
+                    let _ = self.runtime.handle_task_wakeup(*task);
+                    // Running that task, may mark a scope higher up as dirty. If it does, return from the function early
+                    self.process_events();
+                    if self.scopes_need_rerun {
+                        return;
+                    }
+                }
+            }
+
             // Hold a lock to the flush sync to prevent tasks from running in the event we get an immediate
             // When we're doing awaiting the rx, the lock will be dropped and tasks waiting on the lock will get waked
             // We have to own the lock since poll_tasks is cancel safe - the future that this is running in might get dropped
@@ -455,7 +507,12 @@ impl VirtualDom {
 
             match self.rx.next().await.expect("channel should never close") {
                 SchedulerMsg::Immediate(id) => self.mark_dirty(id),
-                SchedulerMsg::TaskNotified(id) => _ = self.runtime.handle_task_wakeup(id),
+                SchedulerMsg::TaskNotified(id) => {
+                    // _ = self.runtime.handle_task_wakeup(id)
+                    // Instead of running the task immediately, we insert it into the runtime's task queue.
+                    // The task may be marked dirty at the same time as the scope that owns the task is dropped.
+                    self.mark_task_dirty(id);
+                }
             };
         }
     }
@@ -468,7 +525,7 @@ impl VirtualDom {
         while let Ok(Some(msg)) = self.rx.try_next() {
             match msg {
                 SchedulerMsg::Immediate(id) => self.mark_dirty(id),
-                SchedulerMsg::TaskNotified(task) => _ = self.runtime.handle_task_wakeup(task),
+                SchedulerMsg::TaskNotified(task) => self.mark_task_dirty(task),
             }
         }
     }
@@ -489,10 +546,8 @@ impl VirtualDom {
                 {
                     let context = scope.state();
                     let height = context.height;
-                    self.dirty_scopes.insert(DirtyScope {
-                        height,
-                        id: context.id,
-                    });
+                    self.dirty_scopes
+                        .insert(DirtyScope::new(height, context.id));
                 }
             }
         }
@@ -558,18 +613,26 @@ impl VirtualDom {
         // We choose not to poll the deadline since we complete pretty quickly anyways
         while let Some(dirty) = self.dirty_scopes.pop_first() {
             // If the scope doesn't exist for whatever reason, then we should skip it
-            if !self.scopes.contains(dirty.id.0) {
+            if !self.scopes.contains(dirty.order.id.0) {
                 continue;
             }
 
             {
                 let _runtime = RuntimeGuard::new(self.runtime.clone());
-                // Run the scope and get the mutations
-                let new_nodes = self.run_scope(dirty.id);
+                // Poll any tasks that might be pending in the scope
+                for task in dirty.tasks_queued.borrow().iter() {
+                    let _ = self.runtime.handle_task_wakeup(*task);
+                }
+                // If the scope is dirty, run the scope and get the mutations
+                if dirty.rerun_queued.get() {
+                    let new_nodes = self.run_scope(dirty.order.id);
 
-                self.diff_scope(to, dirty.id, new_nodes);
+                    self.diff_scope(to, dirty.order.id, new_nodes);
+                }
             }
         }
+
+        self.scopes_need_rerun = false;
     }
 
     /// [`Self::render_immediate`] to a vector of mutations for testing purposes