feat: abstract to rc wake

packages/core/src/scheduler/handle.rs

@@ -16,8 +16,12 @@ impl std::ops::Deref for SchedulerHandle {
 
 pub struct HandleInner {
     pub sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
-    pub tasks: RefCell<Slab<LocalTask>>,
-    pub leaves: RefCell<Slab<SuspenseLeaf>>,
+
+    /// Tasks created with cx.spawn
+    pub tasks: RefCell<Slab<Rc<LocalTask>>>,
+
+    /// Async components
+    pub leaves: RefCell<Slab<Rc<SuspenseLeaf>>>,
 }
 
 impl SchedulerHandle {

packages/core/src/scheduler/mod.rs

@@ -6,6 +6,7 @@ mod handle;
 mod suspense;
 mod task;
 mod wait;
+mod waker;
 
 pub use handle::*;
 use slab::Slab;

packages/core/src/scheduler/task.rs

@@ -1,12 +1,21 @@
-use std::{cell::RefCell, mem, ops::DerefMut, pin::Pin, process::Output, rc::Rc, sync::Arc};
+use std::{
+    cell::{RefCell, UnsafeCell},
+    marker::PhantomData,
+    mem::{self, MaybeUninit},
+    ops::DerefMut,
+    pin::Pin,
+    process::Output,
+    rc::Rc,
+    sync::Arc,
+};
 
 use futures_task::{waker, ArcWake, Context, RawWaker, RawWakerVTable, Waker};
 use futures_util::{pin_mut, Future, FutureExt};
 use slab::Slab;
 
-use crate::ScopeId;
+use crate::{Element, ScopeId};
 
-use super::{HandleInner, SchedulerHandle, SchedulerMsg};
+use super::{waker::RcWake, HandleInner, SchedulerHandle, SchedulerMsg};
 
 #[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
@@ -14,27 +23,25 @@ pub struct TaskId(pub usize);
 
 /// the task itself is the waker
 
-#[derive(Clone)]
 pub struct LocalTask {
     id: TaskId,
     scope: ScopeId,
     tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
-    pub task: *mut dyn Future<Output = ()>,
+    pub task: UnsafeCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
 }
 
 impl HandleInner {
     pub fn spawn(&self, scope: ScopeId, task: impl Future<Output = ()> + 'static) -> TaskId {
         let mut tasks = self.tasks.borrow_mut();
-
         let entry = tasks.vacant_entry();
         let task_id = TaskId(entry.key());
 
-        entry.insert(LocalTask {
+        entry.insert(Rc::new(LocalTask {
             id: task_id,
             tx: self.sender.clone(),
-            task: Box::into_raw(Box::new(task)),
+            task: UnsafeCell::new(Box::pin(task)),
             scope,
-        });
+        }));
 
         self.sender
             .unbounded_send(SchedulerMsg::TaskNotified(task_id))
@@ -54,28 +61,15 @@ impl HandleInner {
     }
 }
 
-pub fn make_task_waker(task: &LocalTask) -> Waker {
-    let raw = RawWaker::new(task as *const LocalTask as *const _, task_vtable());
-    unsafe { Waker::from_raw(raw) }
+pub fn make_task_waker(task: Rc<LocalTask>) -> Waker {
+    let ptr = Rc::into_raw(task).cast::<()>();
+    super::waker::make_rc_waker(task)
 }
 
-fn task_vtable() -> &'static RawWakerVTable {
-    &RawWakerVTable::new(clone, wake, wake_by_ref, drop_task)
-}
-
-unsafe fn clone(data: *const ()) -> RawWaker {
-    RawWaker::new(data as *const (), task_vtable())
-}
-unsafe fn wake(data: *const ()) {
-    wake_by_ref(data);
-}
-unsafe fn wake_by_ref(data: *const ()) {
-    let task = &*(data as *const LocalTask);
-    task.tx
-        .unbounded_send(SchedulerMsg::TaskNotified(task.id))
-        .expect("Scheduler should exist");
-}
-
-unsafe fn drop_task(_data: *const ()) {
-    // doesnt do anything
+impl RcWake for LocalTask {
+    fn wake_by_ref(arc_self: &Rc<Self>) {
+        _ = arc_self
+            .tx
+            .unbounded_send(SchedulerMsg::TaskNotified(arc_self.id));
+    }
 }

packages/core/src/scheduler/wait.rs

@@ -1,7 +1,5 @@
-use std::{ops::DerefMut, pin::Pin};
-
 use futures_task::Context;
-use futures_util::StreamExt;
+use futures_util::{FutureExt, StreamExt};
 
 use crate::{innerlude::make_task_waker, VirtualDom};
 
@@ -13,38 +11,28 @@ impl VirtualDom {
     /// This is cancel safe, so if the future is dropped, you can push events into the virtualdom
     pub async fn wait_for_work(&mut self) {
         loop {
-            let msg = self.scheduler.rx.next().await.unwrap();
-
-            println!("msg received: {:?}", msg);
-
-            match msg {
+            match self.scheduler.rx.next().await.unwrap() {
                 SchedulerMsg::Event => todo!(),
                 SchedulerMsg::Immediate(_) => todo!(),
                 SchedulerMsg::DirtyAll => todo!(),
+
                 SchedulerMsg::TaskNotified(id) => {
                     let mut tasks = self.scheduler.handle.tasks.borrow_mut();
                     let local_task = &tasks[id.0];
 
-                    // // attach the waker to itself
-                    let waker = make_task_waker(local_task);
+                    // attach the waker to itself
+                    // todo: don't make a new waker every time, make it once and then just clone it
+                    let waker = make_task_waker(local_task.clone());
                     let mut cx = Context::from_waker(&waker);
-                    let mut fut = unsafe { &mut *local_task.task };
 
-                    let pinned = unsafe { Pin::new_unchecked(fut.deref_mut()) };
+                    // safety: the waker owns its task and everythig is single threaded
+                    let fut = unsafe { &mut *local_task.task.get() };
 
-                    match pinned.poll(&mut cx) {
-                        futures_task::Poll::Ready(_) => {
-                            // remove the task
-                            tasks.remove(id.0);
-                        }
-                        futures_task::Poll::Pending => {}
-                    }
-
-                    if tasks.is_empty() {
-                        return;
+                    if let futures_task::Poll::Ready(_) = fut.poll_unpin(&mut cx) {
+                        tasks.remove(id.0);
                     }
                 }
-                // SchedulerMsg::TaskNotified(id) => {},
+
                 SchedulerMsg::SuspenseNotified(_) => todo!(),
             }
         }

packages/core/src/scheduler/waker.rs

@@ -0,0 +1,54 @@
+use std::{
+    cell::{RefCell, UnsafeCell},
+    marker::PhantomData,
+    mem::{self, MaybeUninit},
+    ops::DerefMut,
+    pin::Pin,
+    process::Output,
+    rc::Rc,
+    sync::Arc,
+};
+
+use futures_task::{waker, RawWaker, RawWakerVTable, Waker};
+
+pub trait RcWake {
+    fn wake(self: Rc<Self>) {
+        Self::wake_by_ref(&self)
+    }
+    fn wake_by_ref(arc_self: &Rc<Self>);
+}
+
+pub fn make_rc_waker<T: RcWake>(rc: Rc<T>) -> Waker {
+    unsafe { Waker::from_raw(RawWaker::new(Rc::into_raw(rc).cast(), rc_vtable::<T>())) }
+}
+
+fn rc_vtable<T: RcWake>() -> &'static RawWakerVTable {
+    &RawWakerVTable::new(
+        clone_rc_raw::<T>,
+        wake_rc_raw::<T>,
+        wake_by_ref_rc_raw::<T>,
+        drop_rc_raw::<T>,
+    )
+}
+
+// FIXME: panics on Rc::clone / refcount changes could wreak havoc on the
+// code here. We should guard against this by aborting.
+unsafe fn clone_rc_raw<T: RcWake>(data: *const ()) -> RawWaker {
+    let arc = mem::ManuallyDrop::new(Rc::<T>::from_raw(data.cast::<T>()));
+    let _rc_clone: mem::ManuallyDrop<_> = arc.clone();
+    RawWaker::new(data, rc_vtable::<T>())
+}
+
+unsafe fn wake_rc_raw<T: RcWake>(data: *const ()) {
+    let arc: Rc<T> = Rc::from_raw(data.cast::<T>());
+    arc.wake();
+}
+
+unsafe fn wake_by_ref_rc_raw<T: RcWake>(data: *const ()) {
+    let arc = mem::ManuallyDrop::new(Rc::<T>::from_raw(data.cast::<T>()));
+    arc.wake();
+}
+
+unsafe fn drop_rc_raw<T: RcWake>(data: *const ()) {
+    drop(Rc::<T>::from_raw(data.cast::<T>()))
+}