From 55e6dd9701d3710ff17a25581b99b8f159e609b9 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Fri, 12 Nov 2021 01:36:33 -0500
Subject: [PATCH] feat: wire tasks up

---
 packages/core/src/component.rs   |  2 +-
 packages/core/src/scopearena.rs  |  8 ++++---
 packages/core/src/virtual_dom.rs | 39 +++++++++++++++++++++-----------
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/packages/core/src/component.rs b/packages/core/src/component.rs
index 959b17251..9aa5bf2ee 100644
--- a/packages/core/src/component.rs
+++ b/packages/core/src/component.rs
@@ -40,7 +40,7 @@ impl<const A: bool> FragmentBuilder<A> {
 /// fn App(cx: Context, props: &()) -> Element {
 ///     cx.render(rsx!{
 ///         CustomCard {
-///             h1 {}
+///             h1 {}2
 ///             p {}
 ///         }
 ///     })
diff --git a/packages/core/src/scopearena.rs b/packages/core/src/scopearena.rs
index db635352a..315e32a5b 100644
--- a/packages/core/src/scopearena.rs
+++ b/packages/core/src/scopearena.rs
@@ -1,6 +1,6 @@
 use bumpalo::Bump;
 use futures_channel::mpsc::UnboundedSender;
-use fxhash::FxHashMap;
+use fxhash::{FxHashMap, FxHashSet};
 use slab::Slab;
 use std::{
     borrow::Borrow,
@@ -22,8 +22,9 @@ pub struct Heuristic {
 // has an internal heuristics engine to pre-allocate arenas to the right size
 pub(crate) struct ScopeArena {
     bump: Bump,
+    pub pending_futures: FxHashSet<ScopeId>,
     scope_counter: Cell<usize>,
-    scopes: RefCell<FxHashMap<ScopeId, *mut Scope>>,
+    pub scopes: RefCell<FxHashMap<ScopeId, *mut Scope>>,
     pub heuristics: RefCell<FxHashMap<FcSlot, Heuristic>>,
     free_scopes: RefCell<Vec<*mut Scope>>,
     nodes: RefCell<Slab<*const VNode<'static>>>,
@@ -55,6 +56,7 @@ impl ScopeArena {
         Self {
             scope_counter: Cell::new(0),
             bump,
+            pending_futures: FxHashSet::default(),
             scopes: RefCell::new(FxHashMap::default()),
             heuristics: RefCell::new(FxHashMap::default()),
             free_scopes: RefCell::new(Vec::new()),
@@ -354,7 +356,7 @@ impl ScopeArena {
             debug_assert_eq!(scope.wip_frame().nodes.borrow().len(), 1);
 
             if !scope.items.borrow().tasks.is_empty() {
-                // self.
+                //
             }
 
             // make the "wip frame" contents the "finished frame"
diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs
index d014fc3c0..79cb68c06 100644
--- a/packages/core/src/virtual_dom.rs
+++ b/packages/core/src/virtual_dom.rs
@@ -7,6 +7,7 @@ use futures_channel::mpsc::{UnboundedReceiver, UnboundedSender};
 use futures_util::{Future, StreamExt};
 use fxhash::FxHashSet;
 use indexmap::IndexSet;
+use smallvec::SmallVec;
 use std::pin::Pin;
 use std::rc::Rc;
 use std::sync::Arc;
@@ -116,8 +117,6 @@ pub struct VirtualDom {
 
     sender: UnboundedSender<SchedulerMsg>,
 
-    pending_futures: FxHashSet<ScopeId>,
-
     pending_messages: VecDeque<SchedulerMsg>,
 
     dirty_scopes: IndexSet<ScopeId>,
@@ -210,7 +209,6 @@ impl VirtualDom {
             receiver,
             _root_props: caller,
             pending_messages,
-            pending_futures: Default::default(),
             dirty_scopes,
             sender,
         }
@@ -266,45 +264,61 @@ impl VirtualDom {
         // todo: poll the events once even if there is work to do to prevent starvation
 
         // if there's no futures in the virtualdom, just wait for a scheduler message and put it into the queue to be processed
-        if self.pending_futures.is_empty() {
+        if self.scopes.pending_futures.is_empty() {
             self.pending_messages
                 .push_front(self.receiver.next().await.unwrap());
         } else {
             struct PollTasks<'a> {
-                pending_futures: &'a FxHashSet<ScopeId>,
-                scopes: &'a ScopeArena,
+                scopes: &'a mut ScopeArena,
             }
 
             impl<'a> Future for PollTasks<'a> {
                 type Output = ();
 
                 fn poll(
-                    self: Pin<&mut Self>,
+                    mut self: Pin<&mut Self>,
                     cx: &mut std::task::Context<'_>,
                 ) -> Poll<Self::Output> {
                     let mut all_pending = true;
 
+                    let mut unfinished_tasks: SmallVec<[_; 10]> = smallvec::smallvec![];
+                    let mut scopes_to_clear: SmallVec<[_; 10]> = smallvec::smallvec![];
+
                     // Poll every scope manually
-                    for fut in self.pending_futures.iter() {
+                    for fut in self.scopes.pending_futures.iter() {
                         let scope = self
                             .scopes
                             .get_scope(fut)
                             .expect("Scope should never be moved");
 
                         let mut items = scope.items.borrow_mut();
-                        for task in items.tasks.iter_mut() {
-                            let task = task.as_mut();
 
+                        while let Some(mut task) = items.tasks.pop() {
                             // todo: does this make sense?
                             // I don't usually write futures by hand
                             // I think the futures neeed to be pinned using bumpbox or something
                             // right now, they're bump allocated so this shouldn't matter anyway - they're not going to move
-                            let unpinned = unsafe { Pin::new_unchecked(task) };
+                            let task_mut = task.as_mut();
+                            let unpinned = unsafe { Pin::new_unchecked(task_mut) };
 
                             if unpinned.poll(cx).is_ready() {
                                 all_pending = false
+                            } else {
+                                unfinished_tasks.push(task);
                             }
                         }
+
+                        if unfinished_tasks.is_empty() {
+                            scopes_to_clear.push(*fut);
+                        }
+
+                        for task in unfinished_tasks.drain(..) {
+                            items.tasks.push(task);
+                        }
+                    }
+
+                    for scope in scopes_to_clear {
+                        self.scopes.pending_futures.remove(&scope);
                     }
 
                     // Resolve the future if any singular task is ready
@@ -320,8 +334,7 @@ impl VirtualDom {
 
             let scheduler_fut = self.receiver.next();
             let tasks_fut = PollTasks {
-                pending_futures: &self.pending_futures,
-                scopes: &self.scopes,
+                scopes: &mut self.scopes,
             };
 
             match select(tasks_fut, scheduler_fut).await {