clean up dirty scopes module

packages/core/src/virtual_dom.rs

@@ -2,14 +2,14 @@
 //!
 //! This module provides the primary mechanics to create a hook-based, concurrent VDOM for Rust.
 
-use crate::innerlude::ScopeOrder;
+use crate::innerlude::{DirtyTasks, ScopeOrder};
 use crate::Task;
 use crate::{
     any_props::AnyProps,
     arena::ElementId,
     innerlude::{
-        DirtyScopes, ElementRef, ErrorBoundary, NoOpMutations, SchedulerMsg, ScopeState,
-        VNodeMount, VProps, WriteMutations,
+        ElementRef, ErrorBoundary, NoOpMutations, SchedulerMsg, ScopeState, VNodeMount, VProps,
+        WriteMutations,
     },
     nodes::RenderReturn,
     nodes::{Template, TemplateId},
@@ -20,6 +20,7 @@ use crate::{
 use futures_util::StreamExt;
 use rustc_hash::FxHashMap;
 use slab::Slab;
+use std::collections::BTreeSet;
 use std::{any::Any, rc::Rc};
 use tracing::instrument;
 
@@ -185,7 +186,8 @@ use tracing::instrument;
 pub struct VirtualDom {
     pub(crate) scopes: Slab<ScopeState>,
 
-    pub(crate) dirty_scopes: DirtyScopes,
+    pub(crate) dirty_scopes: BTreeSet<ScopeOrder>,
+    pub(crate) dirty_tasks: BTreeSet<DirtyTasks>,
 
     // Maps a template path to a map of byte indexes to templates
     pub(crate) templates: FxHashMap<TemplateId, FxHashMap<usize, Template>>,
@@ -312,6 +314,7 @@ impl VirtualDom {
             runtime: Runtime::new(tx),
             scopes: Default::default(),
             dirty_scopes: Default::default(),
+            dirty_tasks: Default::default(),
             templates: Default::default(),
             queued_templates: Default::default(),
             elements: Default::default(),
@@ -376,7 +379,8 @@ impl VirtualDom {
 
         tracing::event!(tracing::Level::TRACE, "Marking scope {:?} as dirty", id);
         let order = ScopeOrder::new(scope.height(), id);
-        self.dirty_scopes.queue_scope(order);
+        drop(scope);
+        self.queue_scope(order);
     }
 
     /// Mark a task as dirty
@@ -396,7 +400,8 @@ impl VirtualDom {
         );
 
         let order = ScopeOrder::new(scope.height(), scope.id);
-        self.dirty_scopes.queue_task(task, order);
+        drop(scope);
+        self.queue_task(task, order);
     }
 
     /// 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.**
@@ -450,7 +455,7 @@ 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.has_dirty_scopes() {
+            if self.has_dirty_scopes() {
                 return;
             }
 
@@ -491,7 +496,7 @@ impl VirtualDom {
         self.queue_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.has_dirty_scopes() {
+        if self.has_dirty_scopes() {
             return;
         }
 
@@ -505,19 +510,14 @@ impl VirtualDom {
         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(task) = self.dirty_scopes.pop_task() {
-            // If the scope doesn't exist for whatever reason, then we should skip it
-            if !self.scopes.contains(task.order.id.0) {
-                continue;
-            }
-
+        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 _ = 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.dirty_scopes.has_dirty_scopes() {
+                if self.has_dirty_scopes() {
                     return;
                 }
             }
@@ -609,12 +609,7 @@ impl VirtualDom {
 
         // Next, diff any dirty scopes
         // We choose not to poll the deadline since we complete pretty quickly anyways
-        while let Some(work) = self.dirty_scopes.pop_work() {
-            // If the scope doesn't exist for whatever reason, then we should skip it
-            if !self.scopes.contains(work.scope.id.0) {
-                continue;
-            }
-
+        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
@@ -660,7 +655,7 @@ impl VirtualDom {
                 self.queue_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.has_dirty_scopes() {
+                if self.has_dirty_scopes() {
                     break;
                 }
 
@@ -669,12 +664,7 @@ impl VirtualDom {
                     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(task) = self.dirty_scopes.pop_task() {
-                        // If the scope doesn't exist for whatever reason, then we should skip it
-                        if !self.scopes.contains(task.order.id.0) {
-                            continue;
-                        }
-
+                    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 {
@@ -682,7 +672,7 @@ impl VirtualDom {
                                 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.dirty_scopes.has_dirty_scopes() {
+                                if self.has_dirty_scopes() {
                                     break 'wait_for_work;
                                 }
                             }
@@ -694,27 +684,20 @@ impl VirtualDom {
             }
 
             // Render whatever work needs to be rendered, unlocking new futures and suspense leaves
-            while let Some(work) = self.dirty_scopes.pop_work() {
-                // If the scope doesn't exist for whatever reason, then we should skip it
-                if !self.scopes.contains(work.scope.id.0) {
-                    continue;
+            let _runtime = RuntimeGuard::new(self.runtime.clone());
+            while let Some(work) = self.pop_work() {
+                // Then, poll any tasks that might be pending in the scope
+                for task in work.tasks {
+                    // During suspense, we only want to run tasks that are suspended
+                    if self.runtime.suspended_tasks.borrow().contains(&task) {
+                        let _ = self.runtime.handle_task_wakeup(task);
+                    }
                 }
+                // 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);
 
-                {
-                    let _runtime = RuntimeGuard::new(self.runtime.clone());
-                    // Then, poll any tasks that might be pending in the scope
-                    for task in work.tasks {
-                        // During suspense, we only want to run tasks that are suspended
-                        if self.runtime.suspended_tasks.borrow().contains(&task) {
-                            let _ = self.runtime.handle_task_wakeup(task);
-                        }
-                    }
-                    // 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);
-
-                        self.diff_scope(&mut NoOpMutations, work.scope.id, new_nodes);
-                    }
+                    self.diff_scope(&mut NoOpMutations, work.scope.id, new_nodes);
                 }
             }
         }