From 9496853fbcb1d783cacb4c0c226579fa6092877c Mon Sep 17 00:00:00 2001
From: Evan Almloff <evanalmloff@gmail.com>
Date: Fri, 8 Sep 2023 10:05:58 -0500
Subject: [PATCH] expose public methods for manipulating the global runtime

---
 packages/core/src/lib.rs         | 11 ++++---
 packages/core/src/runtime.rs     | 56 +++++++++++++++++++++++++++++---
 packages/core/src/virtual_dom.rs |  5 +++
 3 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/packages/core/src/lib.rs b/packages/core/src/lib.rs
index 62e142f60..5b2f83305 100644
--- a/packages/core/src/lib.rs
+++ b/packages/core/src/lib.rs
@@ -32,6 +32,7 @@ pub(crate) mod innerlude {
     pub use crate::nodes::RenderReturn;
     pub use crate::nodes::*;
     pub use crate::properties::*;
+    pub use crate::runtime::{in_runtime, override_runtime, Runtime};
     pub use crate::scheduler::*;
     pub use crate::scope_context::*;
     pub use crate::scopes::*;
@@ -86,11 +87,11 @@ pub use crate::innerlude::{
 pub mod prelude {
     pub use crate::innerlude::{
         consume_context, consume_context_from_scope, current_scope_id, fc_to_builder, has_context,
-        provide_context, provide_context_to_scope, provide_root_context, push_future,
-        remove_future, schedule_update_any, spawn, spawn_forever, suspend, throw, AnyValue,
-        Component, Element, Event, EventHandler, Fragment, IntoAttributeValue, LazyNodes,
-        Properties, Scope, ScopeId, ScopeState, Scoped, TaskId, Template, TemplateAttribute,
-        TemplateNode, Throw, VNode, VirtualDom,
+        in_runtime, override_runtime, provide_context, provide_context_to_scope,
+        provide_root_context, push_future, remove_future, schedule_update_any, spawn,
+        spawn_forever, suspend, throw, AnyValue, Component, Element, Event, EventHandler, Fragment,
+        IntoAttributeValue, LazyNodes, Properties, Runtime, Scope, ScopeId, ScopeState, Scoped,
+        TaskId, Template, TemplateAttribute, TemplateNode, Throw, VNode, VirtualDom,
     };
 }
 
diff --git a/packages/core/src/runtime.rs b/packages/core/src/runtime.rs
index b2b535830..ea518dbd2 100644
--- a/packages/core/src/runtime.rs
+++ b/packages/core/src/runtime.rs
@@ -7,6 +7,39 @@ thread_local! {
     static RUNTIMES: RefCell<Vec<Rc<Runtime>>> = RefCell::new(vec![]);
 }
 
+/// Run some code within a runtime
+pub fn in_runtime<R>(runtime: Rc<Runtime>, f: impl FnOnce() -> R) -> R {
+    let _guard = RuntimeGuard::new(runtime);
+    f()
+}
+
+/// Override the current runtime. This must be used to override the current runtime when importing components from a dynamic library that has it's own runtime.
+/// 
+/// ```rust
+/// use dioxus::prelude::*;
+/// 
+/// fn main() {
+///     let virtual_dom = VirtualDom::new(app);
+/// }
+/// 
+/// fn app(cx: Scope) -> Element {
+///     render!{ Component { runtime: Runtime::current().unwrap() } }
+/// }
+/// 
+/// // In a dynamic library
+/// #[inline_props]
+/// fn Component(cx: Scope, runtime: Rc<Runtime>) -> Element {
+///     cx.use_hook(|| override_runtime(runtime.clone()));
+/// }
+/// ```
+pub fn override_runtime(runtime: Rc<Runtime>) {
+    RUNTIMES.with(|stack| {
+        let mut stack = stack.borrow_mut();
+        stack.pop();
+        stack.push(runtime);
+    });
+}
+
 /// Pushes a new scope onto the stack
 pub(crate) fn push_runtime(runtime: Rc<Runtime>) {
     RUNTIMES.with(|stack| stack.borrow_mut().push(runtime));
@@ -41,7 +74,8 @@ where
     .flatten()
 }
 
-pub(crate) struct Runtime {
+/// 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) scheduler: Rc<Scheduler>,
 
@@ -63,6 +97,11 @@ impl Runtime {
         })
     }
 
+    /// Get the current runtime
+    pub fn current() -> Option<Rc<Self>> {
+        RUNTIMES.with(|stack| stack.borrow().last().cloned())
+    }
+
     /// 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();
@@ -77,16 +116,25 @@ impl Runtime {
     }
 
     /// Get the current scope id
-    pub fn current_scope_id(&self) -> Option<ScopeId> {
+    pub(crate) fn current_scope_id(&self) -> Option<ScopeId> {
         self.scope_stack.borrow().last().copied()
     }
 
     /// Get the context for any scope given its ID
     ///
     /// This is useful for inserting or removing contexts from a scope, or rendering out its root node
-    pub fn get_context(&self, id: ScopeId) -> Option<Ref<'_, ScopeContext>> {
+    pub(crate) fn get_context(&self, id: ScopeId) -> Option<Ref<'_, ScopeContext>> {
         Ref::filter_map(self.scope_contexts.borrow(), |contexts| {
-            contexts.get(id.0).and_then(|f| f.as_ref())
+            match contexts.get(id.0) {
+                Some(context) => context.as_ref(),
+                _ => {
+                    tracing::error!(
+                        "Attempted to get context for scope that doesn't exist (id: {})",
+                        id.0
+                    );
+                    None
+                }
+            }
         })
         .ok()
     }
diff --git a/packages/core/src/virtual_dom.rs b/packages/core/src/virtual_dom.rs
index a1def71d2..aee0e3d2e 100644
--- a/packages/core/src/virtual_dom.rs
+++ b/packages/core/src/virtual_dom.rs
@@ -658,6 +658,11 @@ impl VirtualDom {
     fn finalize(&mut self) -> Mutations {
         std::mem::take(&mut self.mutations)
     }
+
+    /// Get the current runtime
+    pub fn runtime(&self) -> Rc<Runtime> {
+        self.runtime.clone()
+    }
 }
 
 impl Drop for VirtualDom {