From 9d0d5d74f64a733b4e43ea5cd5cfc7b0c2d2d437 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Tue, 16 Jan 2024 17:52:59 -0800
Subject: [PATCH] Remove last bit of unsafe

---
 packages/core/src/runtime.rs     | 29 ++++++++++-----------
 packages/core/src/virtual_dom.rs | 43 ++++++++++++++------------------
 2 files changed, 34 insertions(+), 38 deletions(-)

diff --git a/packages/core/src/runtime.rs b/packages/core/src/runtime.rs
index 7488f7fe0..5cb8dea48 100644
--- a/packages/core/src/runtime.rs
+++ b/packages/core/src/runtime.rs
@@ -1,14 +1,13 @@
-use std::cell::{Cell, Ref, RefCell};
-
-use slab::Slab;
-
 use crate::{
     innerlude::{LocalTask, SchedulerMsg},
     scope_context::ScopeContext,
     scopes::ScopeId,
     Task,
 };
-use std::rc::Rc;
+use std::{
+    cell::{Cell, Ref, RefCell},
+    rc::Rc,
+};
 
 thread_local! {
     static RUNTIMES: RefCell<Vec<Rc<Runtime>>> = RefCell::new(vec![]);
@@ -27,19 +26,19 @@ pub struct Runtime {
     pub(crate) rendering: Cell<bool>,
 
     /// Tasks created with cx.spawn
-    pub(crate) tasks: RefCell<Slab<LocalTask>>,
+    pub(crate) tasks: RefCell<slab::Slab<LocalTask>>,
 
     pub(crate) sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>,
 }
 
 impl Runtime {
-    pub(crate) fn new(tx: futures_channel::mpsc::UnboundedSender<SchedulerMsg>) -> Rc<Self> {
+    pub(crate) fn new(sender: futures_channel::mpsc::UnboundedSender<SchedulerMsg>) -> Rc<Self> {
         Rc::new(Self {
+            sender,
+            rendering: Cell::new(true),
             scope_contexts: Default::default(),
             scope_stack: Default::default(),
             current_task: Default::default(),
-            rendering: Cell::new(true),
-            sender: tx,
             tasks: Default::default(),
         })
     }
@@ -59,7 +58,9 @@ impl Runtime {
     }
 
     pub(crate) fn remove_context(&self, id: ScopeId) {
-        self.scope_contexts.borrow_mut()[id.0] = None;
+        if let Some(_scope) = self.scope_contexts.borrow_mut()[id.0].take() {
+            // todo: some cleanup work
+        }
     }
 
     /// Get the current scope id
@@ -88,12 +89,12 @@ impl Runtime {
     }
 
     /// Pushes a new scope onto the stack
-    pub(crate) fn push_runtime(runtime: Rc<Runtime>) {
+    pub(crate) fn push(runtime: Rc<Runtime>) {
         RUNTIMES.with(|stack| stack.borrow_mut().push(runtime));
     }
 
     /// Pops a scope off the stack
-    pub(crate) fn pop_runtime() {
+    pub(crate) fn pop() {
         RUNTIMES.with(|stack| stack.borrow_mut().pop());
     }
 
@@ -165,13 +166,13 @@ pub struct RuntimeGuard(());
 impl RuntimeGuard {
     /// Create a new runtime guard that sets the current Dioxus runtime. The runtime will be reset when the guard is dropped
     pub fn new(runtime: Rc<Runtime>) -> Self {
-        Runtime::push_runtime(runtime);
+        Runtime::push(runtime);
         Self(())
     }
 }
 
 impl Drop for RuntimeGuard {
     fn drop(&mut self) {
-        Runtime::pop_runtime();
+        Runtime::pop();
     }
 }
diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs
index e41d6a482..d96ce7b43 100644
--- a/packages/core/src/virtual_dom.rs
+++ b/packages/core/src/virtual_dom.rs
@@ -565,6 +565,21 @@ impl VirtualDom {
         }
     }
 
+    /// Rebuild the virtualdom without handling any of the mutations
+    ///
+    /// This is useful for testing purposes and in cases where you render the output of the virtualdom without
+    /// handling any of its mutations.
+    pub fn rebuild_in_place(&mut self) {
+        self.rebuild(&mut NoOpMutations);
+    }
+
+    /// [`VirtualDom::rebuild`] to a vector of mutations for testing purposes
+    pub fn rebuild_to_vec(&mut self) -> Mutations {
+        let mut mutations = Mutations::default();
+        self.rebuild(&mut mutations);
+        mutations
+    }
+
     /// Performs a *full* rebuild of the virtual dom, returning every edit required to generate the actual dom from scratch.
     ///
     /// The mutations item expects the RealDom's stack to be the root of the application.
@@ -596,21 +611,6 @@ impl VirtualDom {
         to.append_children(ElementId(0), m);
     }
 
-    /// Rebuild the virtualdom without handling any of the mutations
-    ///
-    /// This is useful for testing purposes and in cases where you render the output of the virtualdom without
-    /// handling any of its mutations.
-    pub fn rebuild_in_place(&mut self) {
-        self.rebuild(&mut NoOpMutations);
-    }
-
-    /// [`VirtualDom::rebuild`] to a vector of mutations for testing purposes
-    pub fn rebuild_to_vec(&mut self) -> Mutations {
-        let mut mutations = Mutations::default();
-        self.rebuild(&mut mutations);
-        mutations
-    }
-
     /// Render whatever the VirtualDom has ready as fast as possible without requiring an executor to progress
     /// suspended subtrees.
     pub fn render_immediate(&mut self, to: &mut impl WriteMutations) {
@@ -685,17 +685,12 @@ impl VirtualDom {
                 }
             }
 
-            // Poll the suspense leaves in the meantime
-            let mut work = self.wait_for_work();
+            // Wait until the deadline is ready or we have work if there's no work ready
+            let work = self.wait_for_work();
+            pin_mut!(work);
 
-            // safety: this is okay since we don't touch the original future
-            let pinned = unsafe { std::pin::Pin::new_unchecked(&mut work) };
-
-            // If the deadline is exceded (left) then we should return the mutations we have
             use futures_util::future::{select, Either};
-            if let Either::Left((_, _)) = select(&mut deadline, pinned).await {
-                // release the borrowed
-                drop(work);
+            if let Either::Left((_, _)) = select(&mut deadline, &mut work).await {
                 return;
             }
         }