Store hook state in a mutable scoped-TLS (#1831)

packages/yew-functional/Cargo.toml

@@ -20,3 +20,4 @@ yew-services = { path = "../yew-services" }
 yew = { path = "../yew" }
 yew-functional-macro = { path = "../yew-functional-macro" }
 wasm-bindgen = "0.2.60"
+scoped-tls-hkt = "0.1.2"

packages/yew-functional/src/hooks/mod.rs

@@ -30,13 +30,7 @@ where
     HookUpdate: FnOnce(&mut InternalHookState) -> bool,
 {
     // Extract current hook
-    let (hook, process_message) = CURRENT_HOOK.with(|hook_state_holder| {
-        let hook_state_holder = hook_state_holder.try_borrow_mut();
-        let mut hook_state_holder = hook_state_holder.expect("Nested hooks not supported");
-        let mut hook_state = hook_state_holder
-            .as_mut()
-            .expect("No current hook. Hooks can only be called inside function components");
-
+    let (hook, process_message) = CURRENT_HOOK.with(|hook_state| {
         // Determine which hook position we're at and increment for the next hook
         let hook_pos = hook_state.counter;
         hook_state.counter += 1;
@@ -55,29 +49,26 @@ where
         (hook, hook_state.process_message.clone())
     });
 
+    let hook: Rc<RefCell<InternalHookState>> = hook
+        .downcast()
+        .expect("Incompatible hook type. Hooks must always be called in the same order");
+
     let hook_callback = {
         let hook = hook.clone();
         Box::new(move |update: HookUpdate, post_render| {
             let hook = hook.clone();
             process_message(
                 Box::new(move || {
                     let mut hook = hook.borrow_mut();
-                    let hook = hook.downcast_mut::<InternalHookState>();
-                    let hook = hook.expect(
-                        "Incompatible hook type. Hooks must always be called in the same order",
-                    );
-                    update(hook)
+                    update(&mut hook)
                 }),
                 post_render,
             );
         })
     };
-    let mut hook = hook.borrow_mut();
-    let hook = hook.downcast_mut::<InternalHookState>();
-    let mut hook =
-        hook.expect("Incompatible hook type. Hooks must always be called in the same order");
 
     // Execute the actual hook closure we were given. Let it mutate the hook state and let
     // it create a callback that takes the mutable hook state.
+    let mut hook = hook.borrow_mut();
     hook_runner(&mut hook, hook_callback)
 }

packages/yew-functional/src/lib.rs

@@ -14,6 +14,7 @@
 //!
 //! More details about function components and Hooks can be found on [Yew Docs](https://yew.rs/docs/en/next/concepts/function-components)
 
+use scoped_tls_hkt::scoped_thread_local;
 use std::cell::RefCell;
 use std::rc::Rc;
 use yew::html::AnyScope;
@@ -51,9 +52,7 @@ pub use hooks::*;
 /// ```
 pub use yew_functional_macro::function_component;
 
-thread_local! {
-    static CURRENT_HOOK: RefCell<Option<HookState>> = RefCell::new(None);
-}
+scoped_thread_local!(static mut CURRENT_HOOK: HookState);
 
 type Msg = Box<dyn FnOnce() -> bool>;
 type ProcessMessage = Rc<dyn Fn(Msg, bool)>;
@@ -62,7 +61,7 @@ struct HookState {
     counter: usize,
     scope: AnyScope,
     process_message: ProcessMessage,
-    hooks: Vec<Rc<RefCell<dyn std::any::Any>>>,
+    hooks: Vec<Rc<dyn std::any::Any>>,
     destroy_listeners: Vec<Box<dyn FnOnce()>>,
 }
 
@@ -88,23 +87,18 @@ pub struct FunctionComponent<T: FunctionProvider + 'static> {
     _never: std::marker::PhantomData<T>,
     props: T::TProps,
     link: ComponentLink<Self>,
-    hook_state: RefCell<Option<HookState>>,
+    hook_state: RefCell<HookState>,
     message_queue: MsgQueue,
 }
 
 impl<T> FunctionComponent<T>
 where
     T: FunctionProvider,
 {
-    fn swap_hook_state(&self) {
-        CURRENT_HOOK.with(|previous_hook| {
-            std::mem::swap(
-                &mut *previous_hook
-                    .try_borrow_mut()
-                    .expect("Previous hook still borrowed"),
-                &mut *self.hook_state.borrow_mut(),
-            );
-        });
+    fn with_hook_state<R>(&self, f: impl FnOnce() -> R) -> R {
+        let mut hook_state = self.hook_state.borrow_mut();
+        hook_state.counter = 0;
+        CURRENT_HOOK.set(&mut *hook_state, f)
     }
 }
 
@@ -123,7 +117,7 @@ where
             props,
             link: link.clone(),
             message_queue: message_queue.clone(),
-            hook_state: RefCell::new(Some(HookState {
+            hook_state: RefCell::new(HookState {
                 counter: 0,
                 scope,
                 process_message: Rc::new(move |msg, post_render| {
@@ -135,7 +129,7 @@ where
                 }),
                 hooks: vec![],
                 destroy_listeners: vec![],
-            })),
+            }),
         }
     }
 
@@ -156,34 +150,21 @@ where
     }
 
     fn view(&self) -> Html {
-        // Reset hook
-        self.hook_state
-            .try_borrow_mut()
-            .expect("Unexpected concurrent/nested view call")
-            .as_mut()
-            .unwrap()
-            .counter = 0;
-
-        // Load hook
-        self.swap_hook_state();
-
-        let ret = T::run(&self.props);
-
-        // Restore previous hook
-        self.swap_hook_state();
-
-        ret
+        self.with_hook_state(|| T::run(&self.props))
     }
 
     fn destroy(&mut self) {
-        if let Some(ref mut hook_state) = *self.hook_state.borrow_mut() {
-            for hook in hook_state.destroy_listeners.drain(..) {
-                hook()
-            }
+        let mut hook_state = self.hook_state.borrow_mut();
+        for hook in hook_state.destroy_listeners.drain(..) {
+            hook()
         }
     }
 }
 
 pub(crate) fn get_current_scope() -> Option<AnyScope> {
-    CURRENT_HOOK.with(|cell| cell.borrow().as_ref().map(|state| state.scope.clone()))
+    if CURRENT_HOOK.is_set() {
+        Some(CURRENT_HOOK.with(|state| state.scope.clone()))
+    } else {
+        None
+    }
 }