Make use_callback and Callback bring the runtime with them (#2852)

* Move the runtime along with Callback

packages/core/src/events.rs

@@ -1,4 +1,4 @@
-use crate::{global_context::current_scope_id, properties::SuperFrom, Runtime, ScopeId};
+use crate::{properties::SuperFrom, runtime::RuntimeGuard, Runtime, ScopeId};
 use generational_box::GenerationalBox;
 use std::{cell::RefCell, marker::PhantomData, rc::Rc};
 
@@ -430,7 +430,19 @@ impl<Args: 'static, Ret: 'static> PartialEq for Callback<Args, Ret> {
     }
 }
 
-type ExternalListenerCallback<Args, Ret> = Rc<RefCell<dyn FnMut(Args) -> Ret>>;
+pub(super) struct ExternalListenerCallback<Args, Ret> {
+    callback: Rc<RefCell<dyn FnMut(Args) -> Ret>>,
+    runtime: std::rc::Weak<Runtime>,
+}
+
+impl<Args, Ret> Clone for ExternalListenerCallback<Args, Ret> {
+    fn clone(&self) -> Self {
+        Self {
+            callback: self.callback.clone(),
+            runtime: self.runtime.clone(),
+        }
+    }
+}
 
 impl<Args: 'static, Ret: 'static> Callback<Args, Ret> {
     /// Create a new [`Callback`] from an [`FnMut`]. The callback is owned by the current scope and will be dropped when the scope is dropped.
@@ -439,27 +451,32 @@ impl<Args: 'static, Ret: 'static> Callback<Args, Ret> {
     pub fn new<MaybeAsync: SpawnIfAsync<Marker, Ret>, Marker>(
         mut f: impl FnMut(Args) -> MaybeAsync + 'static,
     ) -> Self {
+        let runtime = Runtime::current().unwrap_or_else(|e| panic!("{}", e));
+        let origin = runtime
+            .current_scope_id()
+            .unwrap_or_else(|e| panic!("{}", e));
         let owner = crate::innerlude::current_owner::<generational_box::UnsyncStorage>();
-        let callback = owner.insert(Some(
-            Rc::new(RefCell::new(move |event: Args| f(event).spawn()))
+        let callback = owner.insert(Some(ExternalListenerCallback {
+            callback: Rc::new(RefCell::new(move |event: Args| f(event).spawn()))
                 as Rc<RefCell<dyn FnMut(Args) -> Ret>>,
-        ));
-        Self {
-            callback,
-            origin: current_scope_id().unwrap_or_else(|e| panic!("{}", e)),
-        }
+            runtime: Rc::downgrade(&runtime),
+        }));
+        Self { callback, origin }
     }
 
     /// Leak a new [`Callback`] that will not be dropped unless it is manually dropped.
     #[track_caller]
     pub fn leak(mut f: impl FnMut(Args) -> Ret + 'static) -> Self {
-        let callback =
-            GenerationalBox::leak(Some(Rc::new(RefCell::new(move |event: Args| f(event)))
-                as Rc<RefCell<dyn FnMut(Args) -> Ret>>));
-        Self {
-            callback,
-            origin: current_scope_id().unwrap_or_else(|e| panic!("{}", e)),
-        }
+        let runtime = Runtime::current().unwrap_or_else(|e| panic!("{}", e));
+        let origin = runtime
+            .current_scope_id()
+            .unwrap_or_else(|e| panic!("{}", e));
+        let callback = GenerationalBox::leak(Some(ExternalListenerCallback {
+            callback: Rc::new(RefCell::new(move |event: Args| f(event).spawn()))
+                as Rc<RefCell<dyn FnMut(Args) -> Ret>>,
+            runtime: Rc::downgrade(&runtime),
+        }));
+        Self { callback, origin }
     }
 
     /// Call this callback with the appropriate argument type
@@ -468,16 +485,15 @@ impl<Args: 'static, Ret: 'static> Callback<Args, Ret> {
     #[track_caller]
     pub fn call(&self, arguments: Args) -> Ret {
         if let Some(callback) = self.callback.read().as_ref() {
-            Runtime::with(|rt| {
-                rt.with_scope_on_stack(self.origin, || {
-                    let value = {
-                        let mut callback = callback.borrow_mut();
-                        callback(arguments)
-                    };
-                    value
-                })
+            let runtime = callback
+                .runtime
+                .upgrade()
+                .expect("Callback was called after the runtime was dropped");
+            let _guard = RuntimeGuard::new(runtime.clone());
+            runtime.with_scope_on_stack(self.origin, || {
+                let mut callback = callback.callback.borrow_mut();
+                callback(arguments)
             })
-            .unwrap_or_else(|e| panic!("{}", e))
         } else {
             panic!("Callback was manually dropped")
         }
@@ -497,17 +513,22 @@ impl<Args: 'static, Ret: 'static> Callback<Args, Ret> {
 
     #[doc(hidden)]
     /// This should only be used by the `rsx!` macro.
-    pub fn __set(&mut self, value: ExternalListenerCallback<Args, Ret>) {
-        self.callback.set(Some(value));
+    pub fn __set(&mut self, value: Rc<RefCell<dyn FnMut(Args) -> Ret>>) {
+        self.callback.set(Some(ExternalListenerCallback {
+            callback: value,
+            runtime: Rc::downgrade(&Runtime::current().unwrap()),
+        }));
     }
 
     #[doc(hidden)]
     /// This should only be used by the `rsx!` macro.
-    pub fn __take(&self) -> ExternalListenerCallback<Args, Ret> {
+    pub fn __take(&self) -> Rc<RefCell<dyn FnMut(Args) -> Ret>> {
         self.callback
             .read()
-            .clone()
+            .as_ref()
             .expect("Callback was manually dropped")
+            .callback
+            .clone()
     }
 }
 

packages/core/src/runtime.rs

@@ -94,6 +94,19 @@ impl Runtime {
             .ok_or(RuntimeError::new())
     }
 
+    /// Wrap a closure so that it always runs in the runtime that is currently active
+    pub fn wrap_closure<'a, I, O>(f: impl Fn(I) -> O + 'a) -> impl Fn(I) -> O + 'a {
+        let current_runtime = Self::current().unwrap();
+        let current_scope = current_runtime.current_scope_id().ok();
+        move |input| match current_scope {
+            Some(scope) => current_runtime.on_scope(scope, || f(input)),
+            None => {
+                let _runtime_guard = RuntimeGuard::new(current_runtime.clone());
+                f(input)
+            }
+        }
+    }
+
     /// Create a scope context. This slab is synchronized with the scope slab.
     pub(crate) fn create_scope(&self, context: Scope) {
         let id = context.id;

packages/desktop/src/hooks.rs

@@ -5,8 +5,8 @@ use crate::{
     ShortcutHandle, ShortcutRegistryError, WryEventHandler,
 };
 use dioxus_core::{
-    prelude::{consume_context, current_scope_id, use_hook_with_cleanup},
-    use_hook,
+    prelude::{consume_context, current_scope_id, use_hook_with_cleanup, RuntimeGuard},
+    use_hook, Runtime,
 };
 
 use dioxus_hooks::use_callback;
@@ -20,10 +20,18 @@ pub fn use_window() -> DesktopContext {
 
 /// Register an event handler that runs when a wry event is processed.
 pub fn use_wry_event_handler(
-    handler: impl FnMut(&Event<UserWindowEvent>, &EventLoopWindowTarget<UserWindowEvent>) + 'static,
+    mut handler: impl FnMut(&Event<UserWindowEvent>, &EventLoopWindowTarget<UserWindowEvent>) + 'static,
 ) -> WryEventHandler {
+    // move the runtime into the event handler closure
+    let runtime = Runtime::current().unwrap();
+
     use_hook_with_cleanup(
-        move || window().create_wry_event_handler(handler),
+        move || {
+            window().create_wry_event_handler(move |event, target| {
+                let _runtime_guard = RuntimeGuard::new(runtime.clone());
+                handler(event, target)
+            })
+        },
         move |handler| handler.remove(),
     )
 }
@@ -37,7 +45,11 @@ pub fn use_wry_event_handler(
 pub fn use_muda_event_handler(
     mut handler: impl FnMut(&muda::MenuEvent) + 'static,
 ) -> WryEventHandler {
+    // move the runtime into the event handler closure
+    let runtime = Runtime::current().unwrap();
+
     use_wry_event_handler(move |event, _| {
+        let _runtime_guard = dioxus_core::prelude::RuntimeGuard::new(runtime.clone());
         if let Event::UserEvent(UserWindowEvent::MudaMenuEvent(event)) = event {
             handler(event);
         }
@@ -50,13 +62,16 @@ pub fn use_muda_event_handler(
 /// if you want to load the asset, and `None` if you want to fallback on the default behavior.
 pub fn use_asset_handler(
     name: &str,
-    handler: impl Fn(AssetRequest, RequestAsyncResponder) + 'static,
+    mut handler: impl FnMut(AssetRequest, RequestAsyncResponder) + 'static,
 ) {
+    // wrap the user's handler in something that keeps it up to date
+    let cb = use_callback(move |(asset, responder)| handler(asset, responder));
+
     use_hook_with_cleanup(
         || {
             crate::window().asset_handlers.register_handler(
                 name.to_string(),
-                Box::new(handler),
+                Box::new(move |asset, responder| cb((asset, responder))),
                 current_scope_id().unwrap(),
             );
 
@@ -73,11 +88,11 @@ pub fn use_global_shortcut(
     accelerator: impl IntoAccelerator,
     mut handler: impl FnMut() + 'static,
 ) -> Result<ShortcutHandle, ShortcutRegistryError> {
-    // wrap the user's handler in something that will carry the scope/runtime with it
+    // wrap the user's handler in something that keeps it up to date
     let cb = use_callback(move |_| handler());
 
     use_hook_with_cleanup(
-        move || window().create_shortcut(accelerator.accelerator(), move || cb.call(())),
+        move || window().create_shortcut(accelerator.accelerator(), move || cb(())),
         |handle| {
             if let Ok(handle) = handle {
                 handle.remove();

packages/hooks/src/use_callback.rs

@@ -1,109 +1,26 @@
-use dioxus_core::prelude::{current_scope_id, use_hook, Runtime};
-use dioxus_signals::CopyValue;
-use dioxus_signals::Writable;
+use std::cell::RefCell;
+use std::rc::Rc;
 
-/// A callback that's always current
-///
-/// Whenever this hook is called the inner callback will be replaced with the new callback but the handle will remain.
+use dioxus_core::prelude::use_hook;
+use dioxus_core::prelude::Callback;
+
+/// Create a callback that's always up to date. Whenever this hook is called the inner callback will be replaced with the new callback but the handle will remain.
 ///
 /// There is *currently* no signal tracking on the Callback so anything reading from it will not be updated.
 ///
 /// This API is in flux and might not remain.
 #[doc = include_str!("../docs/rules_of_hooks.md")]
-pub fn use_callback<T, O>(f: impl FnMut(T) -> O + 'static) -> UseCallback<T, O> {
+pub fn use_callback<T: 'static, O: 'static>(f: impl FnMut(T) -> O + 'static) -> Callback<T, O> {
+    let mut callback = Some(f);
+
     // Create a copyvalue with no contents
     // This copyvalue is generic over F so that it can be sized properly
-    let mut inner = use_hook(|| CopyValue::new(None));
-
-    // Every time this hook is called replace the inner callback with the new callback
-    inner.set(Some(f));
-
-    // And then wrap that callback in a boxed callback so we're blind to the size of the actual callback
-    use_hook(|| {
-        let cur_scope = current_scope_id().unwrap();
-        let rt = Runtime::current().unwrap();
-
-        UseCallback {
-            inner: CopyValue::new(Box::new(move |value| {
-                // run this callback in the context of the scope it was created in.
-                let run_callback =
-                    || inner.with_mut(|f: &mut Option<_>| f.as_mut().unwrap()(value));
-                rt.on_scope(cur_scope, run_callback)
-            })),
-        }
-    })
-}
+    let mut inner = use_hook(|| Callback::new(callback.take().unwrap()));
 
-/// This callback is not generic over a return type so you can hold a bunch of callbacks at once
-///
-/// If you need a callback that returns a value, you can simply wrap the closure you pass in that sets a value in its scope
-pub struct UseCallback<T: 'static, O: 'static + ?Sized> {
-    inner: CopyValue<Box<dyn FnMut(T) -> O>>,
-}
-
-impl<T: 'static, O: 'static + ?Sized> PartialEq for UseCallback<T, O> {
-    fn eq(&self, other: &Self) -> bool {
-        self.inner == other.inner
+    if let Some(callback) = callback.take() {
+        // Every time this hook is called replace the inner callback with the new callback
+        inner.__set(Rc::new(RefCell::new(callback)));
     }
-}
-
-impl<T: 'static, O: 'static + ?Sized> std::fmt::Debug for UseCallback<T, O> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("UseCallback")
-            .field("inner", &self.inner.value())
-            .finish()
-    }
-}
 
-impl<T: 'static, O: 'static + ?Sized> Clone for UseCallback<T, O> {
-    fn clone(&self) -> Self {
-        Self { inner: self.inner }
-    }
-}
-impl<T: 'static, O: 'static> Copy for UseCallback<T, O> {}
-
-impl<T, O> UseCallback<T, O> {
-    /// Call the callback
-    pub fn call(&self, value: T) -> O {
-        (self.inner.write_unchecked())(value)
-    }
-}
-
-// This makes UseCallback callable like a normal function
-impl<T, O> std::ops::Deref for UseCallback<T, O> {
-    type Target = dyn Fn(T) -> O;
-
-    fn deref(&self) -> &Self::Target {
-        use std::mem::MaybeUninit;
-
-        // https://github.com/dtolnay/case-studies/tree/master/callable-types
-
-        // First we create a closure that captures something with the Same in memory layout as Self (MaybeUninit<Self>).
-        let uninit_callable = MaybeUninit::<Self>::uninit();
-        // Then move that value into the closure. We assume that the closure now has a in memory layout of Self.
-        let uninit_closure = move |value| Self::call(unsafe { &*uninit_callable.as_ptr() }, value);
-
-        // Check that the size of the closure is the same as the size of Self in case the compiler changed the layout of the closure.
-        let size_of_closure = std::mem::size_of_val(&uninit_closure);
-        assert_eq!(size_of_closure, std::mem::size_of::<Self>());
-
-        // Then cast the lifetime of the closure to the lifetime of &self.
-        fn cast_lifetime<'a, T>(_a: &T, b: &'a T) -> &'a T {
-            b
-        }
-        let reference_to_closure = cast_lifetime(
-            {
-                // The real closure that we will never use.
-                &uninit_closure
-            },
-            #[allow(clippy::missing_transmute_annotations)]
-            // We transmute self into a reference to the closure. This is safe because we know that the closure has the same memory layout as Self so &Closure == &Self.
-            unsafe {
-                std::mem::transmute(self)
-            },
-        );
-
-        // Cast the closure to a trait object.
-        reference_to_closure as &_
-    }
+    inner
 }

packages/hooks/src/use_future.rs

@@ -1,5 +1,5 @@
 #![allow(missing_docs)]
-use crate::{use_callback, use_hook_did_run, use_signal, UseCallback};
+use crate::{use_callback, use_hook_did_run, use_signal};
 use dioxus_core::prelude::*;
 use dioxus_signals::*;
 use std::future::Future;
@@ -82,7 +82,7 @@ where
 pub struct UseFuture {
     task: CopyValue<Task>,
     state: Signal<UseFutureState>,
-    callback: UseCallback<(), Task>,
+    callback: Callback<(), Task>,
 }
 
 /// A signal that represents the state of a future

packages/hooks/src/use_resource.rs

@@ -1,6 +1,6 @@
 #![allow(missing_docs)]
 
-use crate::{use_callback, use_signal, UseCallback};
+use crate::{use_callback, use_signal};
 use dioxus_core::prelude::*;
 use dioxus_signals::*;
 use futures_util::{future, pin_mut, FutureExt, StreamExt};
@@ -110,7 +110,7 @@ pub struct Resource<T: 'static> {
     value: Signal<Option<T>>,
     task: Signal<Task>,
     state: Signal<UseResourceState>,
-    callback: UseCallback<(), Task>,
+    callback: Callback<(), Task>,
 }
 
 impl<T> PartialEq for Resource<T> {