chore: clean up scheduler code

packages/core/src/scope_arena.rs

@@ -1,5 +1,14 @@
-use crate::{innerlude::SuspenseContext, scheduler::RcWake};
-use futures_util::{pin_mut, task::noop_waker_ref};
+use crate::{
+    any_props::AnyProps,
+    arena::ElementId,
+    bump_frame::BumpFrame,
+    factory::RenderReturn,
+    innerlude::{SuspenseId, SuspenseLeaf},
+    scheduler::RcWake,
+    scopes::{ScopeId, ScopeState},
+    virtual_dom::VirtualDom,
+};
+use futures_util::FutureExt;
 use std::{
     mem,
     pin::Pin,
@@ -7,18 +16,8 @@ use std::{
     task::{Context, Poll},
 };
 
-use crate::{
-    any_props::AnyProps,
-    arena::ElementId,
-    bump_frame::BumpFrame,
-    factory::RenderReturn,
-    innerlude::{SuspenseId, SuspenseLeaf},
-    scopes::{ScopeId, ScopeState},
-    virtual_dom::VirtualDom,
-};
-
 impl VirtualDom {
-    pub fn new_scope(&mut self, props: *mut dyn AnyProps<'static>) -> ScopeId {
+    pub fn new_scope(&mut self, props: *mut dyn AnyProps<'static>) -> &mut ScopeState {
         let parent = self.acquire_current_scope_raw();
         let container = self.acquire_current_container();
         let entry = self.scopes.vacant_entry();
@@ -34,18 +33,17 @@ impl VirtualDom {
             placeholder: None.into(),
             node_arena_1: BumpFrame::new(50),
             node_arena_2: BumpFrame::new(50),
+            spawned_tasks: Default::default(),
             render_cnt: Default::default(),
             hook_arena: Default::default(),
             hook_vals: Default::default(),
             hook_idx: Default::default(),
             shared_contexts: Default::default(),
-            tasks: self.scheduler.handle.clone(),
-        });
-
-        id
+            tasks: self.scheduler.clone(),
+        })
     }
 
-    pub fn acquire_current_container(&self) -> ElementId {
+    fn acquire_current_container(&self) -> ElementId {
         self.element_stack
             .last()
             .copied()
@@ -59,39 +57,45 @@ impl VirtualDom {
             .and_then(|id| self.scopes.get_mut(id.0).map(|f| f as *mut ScopeState))
     }
 
-    pub fn run_scope(&mut self, scope_id: ScopeId) -> &mut RenderReturn {
+    pub(crate) unsafe fn run_scope_extend<'a>(
+        &mut self,
+        scope_id: ScopeId,
+    ) -> &'a RenderReturn<'a> {
+        unsafe { self.run_scope(scope_id).extend_lifetime_ref() }
+    }
+
+    pub(crate) fn run_scope(&mut self, scope_id: ScopeId) -> &RenderReturn {
         let mut new_nodes = unsafe {
             let scope = &mut self.scopes[scope_id.0];
             scope.hook_idx.set(0);
 
+            // safety: due to how we traverse the tree, we know that the scope is not currently aliased
             let props: &mut dyn AnyProps = mem::transmute(&mut *scope.props);
-            let res: RenderReturn = props.render(scope);
-            let res: RenderReturn<'static> = mem::transmute(res);
-            res
+            props.render(scope).extend_lifetime()
         };
 
         // immediately resolve futures that can be resolved
         if let RenderReturn::Async(task) = &mut new_nodes {
-            use futures_util::FutureExt;
+            let mut leaves = self.scheduler.leaves.borrow_mut();
 
-            let mut leaves = self.scheduler.handle.leaves.borrow_mut();
             let entry = leaves.vacant_entry();
-            let key = entry.key();
-            let suspense_id = SuspenseId(key);
+            let suspense_id = SuspenseId(entry.key());
 
             let leaf = Rc::new(SuspenseLeaf {
                 scope_id,
                 task: task.as_mut(),
                 id: suspense_id,
-                tx: self.scheduler.handle.sender.clone(),
-                notified: false.into(),
+                tx: self.scheduler.sender.clone(),
+                notified: Default::default(),
             });
 
-            let _leaf = leaf.clone();
             let waker = leaf.waker();
             let mut cx = Context::from_waker(&waker);
+
+            // safety: the task is already pinned in the bump arena
             let mut pinned = unsafe { Pin::new_unchecked(task.as_mut()) };
 
+            // Keep polling until either we get a value or the future is not ready
             loop {
                 match pinned.poll_unpin(&mut cx) {
                     // If nodes are produced, then set it and we can break
@@ -103,8 +107,8 @@ impl VirtualDom {
                     // If no nodes are produced but the future woke up immediately, then try polling it again
                     // This circumvents things like yield_now, but is important is important when rendering
                     // components that are just a stream of immediately ready futures
-                    _ if _leaf.notified.get() => {
-                        _leaf.notified.set(false);
+                    _ if leaf.notified.get() => {
+                        leaf.notified.set(false);
                         continue;
                     }
 
@@ -112,7 +116,7 @@ impl VirtualDom {
                     // Insert the future into fiber leaves and break
                     _ => {
                         entry.insert(leaf);
-                        self.waiting_on.push(suspense_id);
+                        self.collected_leaves.push(suspense_id);
                         break;
                     }
                 };