From 26f5fb80cb1a61d8f5e8e9be2c8e8e8d22912aa3 Mon Sep 17 00:00:00 2001
From: Evan Almloff <evanalmloff@gmail.com>
Date: Fri, 1 Mar 2024 12:18:42 -0600
Subject: [PATCH] clean up the task queue

---
 examples/readme.rs               | 34 ++++++++++++++++++++---
 packages/core/src/virtual_dom.rs | 46 ++++++++++++++++++--------------
 2 files changed, 56 insertions(+), 24 deletions(-)

diff --git a/examples/readme.rs b/examples/readme.rs
index 7c5240086..be4128a99 100644
--- a/examples/readme.rs
+++ b/examples/readme.rs
@@ -12,11 +12,37 @@ fn main() {
 }
 
 fn app() -> Element {
-    let mut count = use_signal(|| 0);
+    let mut vec = use_signal(|| vec![1, 2, 3]);
+
+    let len = vec.len();
 
     rsx! {
-        h1 { "High-Five counter: {count}" }
-        button { onclick: move |_| count += 1, "Up high!" }
-        button { onclick: move |_| count -= 1, "Down low!" }
+        button {
+            onclick: move |_| {
+                let mut vec = vec.write();
+                vec.push(len);
+            },
+            "Add"
+        }
+        button {
+            onclick: move |_| {
+                vec.pop();
+            },
+            "Remove"
+        }
+        for i in 0..len {
+            Child {
+                index: i,
+                vec,
+            }
+        }
+    }
+}
+
+#[component]
+fn Child(index: usize, vec: Signal<Vec<usize>>) -> Element {
+    let item = use_memo(move || vec.read()[index]);
+    rsx! {
+        div { "Item: {item}" }
     }
 }
diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs
index 6541a70e0..88a404412 100644
--- a/packages/core/src/virtual_dom.rs
+++ b/packages/core/src/virtual_dom.rs
@@ -479,25 +479,6 @@ impl VirtualDom {
             // 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
@@ -528,6 +509,30 @@ impl VirtualDom {
                 SchedulerMsg::TaskNotified(task) => self.mark_task_dirty(task),
             }
         }
+
+        // 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.scopes_need_rerun {
+            return;
+        }
+
+        // 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;
+                }
+            }
+        }
     }
 
     /// Replace a template at runtime. This will re-render all components that use this template.
@@ -619,7 +624,8 @@ impl VirtualDom {
 
             {
                 let _runtime = RuntimeGuard::new(self.runtime.clone());
-                // Poll any tasks that might be pending in the scope
+                // Then, poll any tasks that might be pending in the scope
+                // This will run effects, so this **must** be done after the scope is diffed
                 for task in dirty.tasks_queued.borrow().iter() {
                     let _ = self.runtime.handle_task_wakeup(*task);
                 }