From 821a650f775526847158f3109a2b2b55aa3aa6bc Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Thu, 4 Apr 2024 11:46:31 -0700
Subject: [PATCH] Fix: requeue events, process events  (#2236)

* Fix: requeue events, process events

* Fix test: wait_for_work shouldn't wait 3 times when one update happens while rendering
---
 packages/core/src/virtual_dom.rs | 16 ++++++++++++++--
 packages/hooks/tests/memo.rs     |  9 +++------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs
index 0a2144668..d15ae06d6 100644
--- a/packages/core/src/virtual_dom.rs
+++ b/packages/core/src/virtual_dom.rs
@@ -512,12 +512,17 @@ impl VirtualDom {
         // We choose not to poll the deadline since we complete pretty quickly anyways
         while let Some(task) = self.pop_task() {
             // Then poll any tasks that might be pending
-            let tasks = task.tasks_queued.into_inner();
-            for task in tasks {
+            let mut tasks = task.tasks_queued.into_inner();
+            while let Some(task) = tasks.pop() {
                 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.queue_events();
                 if self.has_dirty_scopes() {
+                    // requeue any remaining tasks
+                    for task in tasks {
+                        self.mark_task_dirty(task);
+                    }
                     return;
                 }
             }
@@ -633,10 +638,14 @@ impl VirtualDom {
         while let Some(work) = self.pop_work() {
             {
                 let _runtime = RuntimeGuard::new(self.runtime.clone());
+
                 // Then, poll any tasks that might be pending in the scope
                 for task in work.tasks {
                     let _ = self.runtime.handle_task_wakeup(task);
                 }
+
+                self.queue_events();
+
                 // If the scope is dirty, run the scope and get the mutations
                 if work.rerun_scope {
                     let new_nodes = self.run_scope(work.scope.id);
@@ -718,6 +727,9 @@ impl VirtualDom {
                         let _ = self.runtime.handle_task_wakeup(task);
                     }
                 }
+
+                self.queue_events();
+
                 // If the scope is dirty, run the scope and get the mutations
                 if work.rerun_scope {
                     let new_nodes = self.run_scope(work.scope.id);
diff --git a/packages/hooks/tests/memo.rs b/packages/hooks/tests/memo.rs
index 186370f79..77e59f3d5 100644
--- a/packages/hooks/tests/memo.rs
+++ b/packages/hooks/tests/memo.rs
@@ -29,15 +29,13 @@ async fn memo_updates() {
         if generation() < 2 {
             vec.push(len);
         }
+
         // The memo should always be up to date
         assert_eq!(vec.len(), len_memo());
 
         rsx! {
             for i in 0..len {
-                Child {
-                    index: i,
-                    vec,
-                }
+                Child { index: i, vec }
             }
         }
     }
@@ -57,10 +55,9 @@ async fn memo_updates() {
     dom.rebuild_in_place();
     let mut signal = VEC_SIGNAL.with(|cell| (*cell.borrow()).unwrap());
     // Wait for the signal to update
-    for _ in 0..3 {
+    for _ in 0..2 {
         dom.wait_for_work().await;
         dom.render_immediate(&mut dioxus::dioxus_core::NoOpMutations);
-        println!("Signal: {signal:?}");
     }
     assert_eq!(signal(), vec![0, 1, 2, 3, 4, 5]);
     // Remove each element from the vec