Add flush_sync and pre_render methods to core

packages/core/src/runtime.rs

@@ -1,12 +1,13 @@
+use rustc_hash::FxHashSet;
+
 use crate::{
     innerlude::{LocalTask, SchedulerMsg},
-    scope_context::ScopeContext,
+    scope_context::Scope,
     scopes::ScopeId,
     Task,
 };
 use std::{
     cell::{Cell, Ref, RefCell},
-    collections::VecDeque,
     rc::Rc,
 };
 
@@ -16,7 +17,7 @@ thread_local! {
 
 /// A global runtime that is shared across all scopes that provides the async runtime and context API
 pub struct Runtime {
-    pub(crate) scope_contexts: RefCell<Vec<Option<ScopeContext>>>,
+    pub(crate) scope_states: RefCell<Vec<Option<Scope>>>,
 
     // We use this to track the current scope
     pub(crate) scope_stack: RefCell<Vec<ScopeId>>,
@@ -29,10 +30,10 @@ pub struct Runtime {
     /// Tasks created with cx.spawn
     pub(crate) tasks: RefCell<slab::Slab<Rc<LocalTask>>>,
 
-    /// Queued tasks that are waiting to be polled
-    pub(crate) queued_tasks: Rc<RefCell<VecDeque<Task>>>,
-
     pub(crate) sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
+
+    // Tasks waiting to be manually resumed when we call wait_for_work
+    pub(crate) flush_table: RefCell<FxHashSet<Task>>,
 }
 
 impl Runtime {
@@ -40,11 +41,11 @@ impl Runtime {
         Rc::new(Self {
             sender,
             rendering: Cell::new(true),
-            scope_contexts: Default::default(),
+            scope_states: Default::default(),
             scope_stack: Default::default(),
             current_task: Default::default(),
             tasks: Default::default(),
-            queued_tasks: Rc::new(RefCell::new(VecDeque::new())),
+            flush_table: Default::default(),
         })
     }
 
@@ -54,35 +55,39 @@ impl Runtime {
     }
 
     /// Create a scope context. This slab is synchronized with the scope slab.
-    pub(crate) fn create_context_at(&self, id: ScopeId, context: ScopeContext) {
-        let mut contexts = self.scope_contexts.borrow_mut();
-        if contexts.len() <= id.0 {
-            contexts.resize_with(id.0 + 1, Default::default);
+    pub(crate) fn create_scope(&self, context: Scope) {
+        let id = context.id;
+        let mut scopes = self.scope_states.borrow_mut();
+        if scopes.len() <= id.0 {
+            scopes.resize_with(id.0 + 1, Default::default);
         }
-        contexts[id.0] = Some(context);
+        scopes[id.0] = Some(context);
     }
 
-    pub(crate) fn remove_context(self: &Rc<Self>, id: ScopeId) {
+    pub(crate) fn remove_scope(self: &Rc<Self>, id: ScopeId) {
         {
-            let borrow = self.scope_contexts.borrow();
+            let borrow = self.scope_states.borrow();
             if let Some(scope) = &borrow[id.0] {
                 let _runtime_guard = RuntimeGuard::new(self.clone());
                 // Manually drop tasks, hooks, and contexts inside of the runtime
                 self.on_scope(id, || {
-                    // Drop all spawned tasks
+                    // Drop all spawned tasks - order doesn't matter since tasks don't rely on eachother
+                    // In theory nested tasks might not like this
                     for id in scope.spawned_tasks.take() {
                         self.remove_task(id);
                     }
 
-                    // Drop all hooks
-                    scope.hooks.take();
+                    // Drop all hooks in reverse order in case a hook depends on another hook.
+                    for hook in scope.hooks.take().drain(..).rev() {
+                        drop(hook);
+                    }
 
                     // Drop all contexts
                     scope.shared_contexts.take();
                 });
             }
         }
-        self.scope_contexts.borrow_mut()[id.0].take();
+        self.scope_states.borrow_mut()[id.0].take();
     }
 
     /// Get the current scope id
@@ -104,11 +109,11 @@ impl Runtime {
         o
     }
 
-    /// Get the context for any scope given its ID
+    /// Get the state for any scope given its ID
     ///
     /// This is useful for inserting or removing contexts from a scope, or rendering out its root node
-    pub(crate) fn get_context(&self, id: ScopeId) -> Option<Ref<'_, ScopeContext>> {
-        Ref::filter_map(self.scope_contexts.borrow(), |contexts| {
+    pub(crate) fn get_state(&self, id: ScopeId) -> Option<Ref<'_, Scope>> {
+        Ref::filter_map(self.scope_states.borrow(), |contexts| {
             contexts.get(id.0).and_then(|f| f.as_ref())
         })
         .ok()
@@ -125,34 +130,22 @@ impl Runtime {
     }
 
     /// Runs a function with the current runtime
-    pub(crate) fn with<F, R>(f: F) -> Option<R>
-    where
-        F: FnOnce(&Runtime) -> R,
-    {
-        RUNTIMES.with(|stack| {
-            let stack = stack.borrow();
-            stack.last().map(|r| f(r))
-        })
+    pub(crate) fn with<R>(f: impl FnOnce(&Runtime) -> R) -> Option<R> {
+        RUNTIMES.with(|stack| stack.borrow().last().map(|r| f(r)))
     }
 
     /// Runs a function with the current scope
-    pub(crate) fn with_current_scope<F, R>(f: F) -> Option<R>
-    where
-        F: FnOnce(&ScopeContext) -> R,
-    {
+    pub(crate) fn with_current_scope<R>(f: impl FnOnce(&Scope) -> R) -> Option<R> {
         Self::with(|rt| {
             rt.current_scope_id()
-                .and_then(|scope| rt.get_context(scope).map(|sc| f(&sc)))
+                .and_then(|scope| rt.get_state(scope).map(|sc| f(&sc)))
         })
         .flatten()
     }
 
     /// Runs a function with the current scope
-    pub(crate) fn with_scope<F, R>(scope: ScopeId, f: F) -> Option<R>
-    where
-        F: FnOnce(&ScopeContext) -> R,
-    {
-        Self::with(|rt| rt.get_context(scope).map(|sc| f(&sc))).flatten()
+    pub(crate) fn with_scope<R>(scope: ScopeId, f: impl FnOnce(&Scope) -> R) -> Option<R> {
+        Self::with(|rt| rt.get_state(scope).map(|sc| f(&sc))).flatten()
     }
 }
 
@@ -182,7 +175,9 @@ impl Runtime {
 /// }
 ///
 /// fn Component(cx: ComponentProps) -> Element {
-///     cx.use_hook(|| RuntimeGuard::new(cx.runtime.clone()));
+///     use_hook(|| {
+///         let _guard = RuntimeGuard::new(cx.runtime.clone());
+///     });
 ///
 ///     rsx! { div {} }
 /// }