Add support for use_context hook (#1249)

* Replace mounted with rendered lifecycle method

* fix checks

* fix doc test

* use_context test

* made test yet more sensitive to errors

* fix some merge oddities

* just some style adjustments

* proof of concept use_context hook using scope

* store scope in hook state

* add a test with multiple context types

* subscribe to context changes

* make clippy happy

* reuse dead slots for subscriptions

* clean up and improve test

Co-authored-by: Justin Starry <justin.starry@icloud.com>
Co-authored-by: Mathis <mk@thepeaklab.com>

.gitignore

@@ -3,4 +3,5 @@ target
 Cargo.lock
 orig.*
 /.idea
+/.vscode
 /cmake-build-debug

yew-functional/src/lib.rs

@@ -2,8 +2,12 @@ use std::borrow::Borrow;
 use std::cell::RefCell;
 use std::ops::DerefMut;
 use std::rc::Rc;
+use yew::html::AnyScope;
 use yew::{Component, ComponentLink, Html, Properties};
 
+mod use_context_hook;
+pub use use_context_hook::*;
+
 thread_local! {
     static CURRENT_HOOK: RefCell<Option<HookState>> = RefCell::new(None);
 }
@@ -15,6 +19,7 @@ pub trait Hook {
 type ProcessMessage = Rc<dyn Fn(Box<dyn FnOnce() -> bool>)>;
 struct HookState {
     counter: usize,
+    scope: AnyScope,
     process_message: ProcessMessage,
     hooks: Vec<Rc<RefCell<dyn std::any::Any>>>,
     destroy_listeners: Vec<Box<dyn FnOnce()>>,
@@ -31,6 +36,23 @@ pub struct FunctionComponent<T: FunctionProvider> {
     hook_state: RefCell<Option<HookState>>,
 }
 
+impl<T> FunctionComponent<T>
+where
+    T: FunctionProvider,
+{
+    fn swap_hook_state(&self) {
+        CURRENT_HOOK.with(|previous_hook| {
+            std::mem::swap(
+                previous_hook
+                    .try_borrow_mut()
+                    .expect("Previous hook still borrowed")
+                    .deref_mut(),
+                self.hook_state.borrow_mut().deref_mut(),
+            );
+        });
+    }
+}
+
 impl<T: 'static> Component for FunctionComponent<T>
 where
     T: FunctionProvider,
@@ -39,11 +61,13 @@ where
     type Properties = T::TProps;
 
     fn create(props: Self::Properties, link: ComponentLink<Self>) -> Self {
+        let scope = AnyScope::from(link.clone());
         FunctionComponent {
             _never: std::marker::PhantomData::default(),
             props,
             hook_state: RefCell::new(Some(HookState {
                 counter: 0,
+                scope,
                 process_message: Rc::new(move |msg| link.send_message(msg)),
                 hooks: vec![],
                 destroy_listeners: vec![],
@@ -61,7 +85,6 @@ where
         props != self.props
     }
 
-    //noinspection DuplicatedCode
     fn view(&self) -> Html {
         // Reset hook
         self.hook_state
@@ -70,29 +93,14 @@ where
             .as_mut()
             .unwrap()
             .counter = 0;
+
         // Load hook
-        CURRENT_HOOK.with(|previous_hook| {
-            std::mem::swap(
-                previous_hook
-                    .try_borrow_mut()
-                    .expect("Previous hook still borrowed")
-                    .deref_mut(),
-                self.hook_state.borrow_mut().deref_mut(),
-            );
-        });
+        self.swap_hook_state();
 
         let ret = T::run(&self.props);
 
-        // Unload hook
-        CURRENT_HOOK.with(|previous_hook| {
-            std::mem::swap(
-                previous_hook
-                    .try_borrow_mut()
-                    .expect("Previous hook still borrowed")
-                    .deref_mut(),
-                self.hook_state.borrow_mut().deref_mut(),
-            );
-        });
+        // Restore previous hook
+        self.swap_hook_state();
 
         ret
     }
@@ -347,3 +355,7 @@ where
     // it create a callback that takes the mutable hook state.
     hook_runner(&mut hook, trigger)
 }
+
+pub(crate) fn get_current_scope() -> Option<AnyScope> {
+    CURRENT_HOOK.with(|cell| cell.borrow().as_ref().map(|state| state.scope.clone()))
+}

yew-functional/src/use_context_hook.rs

@@ -0,0 +1,156 @@
+// Naming this file use_context could be confusing. Not least to the IDE.
+use super::{get_current_scope, use_hook, Hook};
+use std::any::TypeId;
+use std::cell::RefCell;
+use std::rc::{Rc, Weak};
+use std::{iter, mem};
+use yew::html::{AnyScope, Renderable, Scope};
+use yew::{Children, Component, ComponentLink, Html, Properties};
+
+type ConsumerCallback<T> = Box<dyn Fn(Rc<T>)>;
+
+#[derive(Clone, PartialEq, Properties)]
+pub struct ContextProviderProps<T: Clone> {
+    pub context: T,
+    pub children: Children,
+}
+
+pub struct ContextProvider<T: Clone + 'static> {
+    context: Rc<T>,
+    children: Children,
+    consumers: RefCell<Vec<Weak<ConsumerCallback<T>>>>,
+}
+
+impl<T: Clone> ContextProvider<T> {
+    /// Add the callback to the subscriber list to be called whenever the context changes.
+    /// The consumer is unsubscribed as soon as the callback is dropped.
+    fn subscribe_consumer(&self, mut callback: Weak<ConsumerCallback<T>>) {
+        let mut consumers = self.consumers.borrow_mut();
+        // consumers re-subscribe on every render. Try to keep the subscriber list small by reusing dead slots.
+        for cb in consumers.iter_mut() {
+            if cb.strong_count() == 0 {
+                mem::swap(cb, &mut callback);
+                return;
+            }
+        }
+
+        // no slot to reuse, this is a new consumer
+        consumers.push(callback);
+    }
+
+    /// Notify all subscribed consumers and remove dropped consumers from the list.
+    fn notify_consumers(&mut self) {
+        let context = &self.context;
+        self.consumers.borrow_mut().retain(|cb| {
+            if let Some(cb) = cb.upgrade() {
+                cb(Rc::clone(&context));
+                true
+            } else {
+                false
+            }
+        });
+    }
+}
+
+impl<T: Clone + 'static> Component for ContextProvider<T> {
+    type Message = ();
+    type Properties = ContextProviderProps<T>;
+
+    fn create(props: Self::Properties, _link: ComponentLink<Self>) -> Self {
+        ContextProvider {
+            children: props.children,
+            context: Rc::new(props.context),
+            consumers: RefCell::new(Vec::new()),
+        }
+    }
+
+    fn update(&mut self, _msg: Self::Message) -> bool {
+        true
+    }
+
+    fn change(&mut self, props: Self::Properties) -> bool {
+        let should_render = if self.children != props.children {
+            self.children = props.children;
+            true
+        } else {
+            false
+        };
+
+        self.context = Rc::new(props.context);
+        self.notify_consumers();
+
+        should_render
+    }
+
+    fn view(&self) -> Html {
+        self.children.render()
+    }
+}
+
+fn find_context_provider_scope<T: 'static + Clone>(
+    scope: &AnyScope,
+) -> Option<Scope<ContextProvider<T>>> {
+    let expected_type_id = TypeId::of::<ContextProvider<T>>();
+    iter::successors(Some(scope), |scope| scope.get_parent())
+        .filter(|scope| scope.get_type_id() == &expected_type_id)
+        .cloned()
+        .map(AnyScope::downcast::<ContextProvider<T>>)
+        .next()
+}
+
+fn with_provider_component<T, F, R>(
+    provider_scope: &Option<Scope<ContextProvider<T>>>,
+    f: F,
+) -> Option<R>
+where
+    T: Clone,
+    F: FnOnce(&ContextProvider<T>) -> R,
+{
+    provider_scope
+        .as_ref()
+        .and_then(|scope| scope.get_component().map(|comp| f(&*comp)))
+}
+
+pub fn use_context<T: 'static + Clone>() -> Option<Rc<T>> {
+    let scope = get_current_scope()
+        .expect("No current Scope. `use_context` can only be called inside functional components");
+
+    struct UseContextState<T2: 'static + Clone> {
+        provider_scope: Option<Scope<ContextProvider<T2>>>,
+        current_context: Option<Rc<T2>>,
+        callback: Option<Rc<ConsumerCallback<T2>>>,
+    }
+    impl<T: 'static + Clone> Hook for UseContextState<T> {
+        fn tear_down(&mut self) {
+            if let Some(cb) = self.callback.take() {
+                drop(cb);
+            }
+        }
+    }
+    use_hook(
+        |state: &mut UseContextState<T>, hook_update| {
+            state.callback = Some(Rc::new(Box::new(move |ctx: Rc<T>| {
+                hook_update(|state: &mut UseContextState<T>| {
+                    state.current_context = Some(ctx);
+                    true
+                });
+            })));
+            let weak_cb = Rc::downgrade(state.callback.as_ref().unwrap());
+            with_provider_component(&state.provider_scope, |comp| {
+                comp.subscribe_consumer(weak_cb)
+            });
+
+            state.current_context.clone()
+        },
+        move || {
+            let provider_scope = find_context_provider_scope::<T>(&scope);
+            let current_context =
+                with_provider_component(&provider_scope, |comp| Rc::clone(&comp.context));
+            UseContextState {
+                provider_scope,
+                current_context,
+                callback: None,
+            }
+        },
+    )
+}

yew-functional/tests/lib.rs

@@ -436,8 +436,8 @@ fn use_effect_refires_on_dependency_change() {
 }
 
 fn obtain_result() -> String {
-    return yew::utils::document()
+    yew::utils::document()
         .get_element_by_id("result")
         .expect("No result found. Most likely, the application crashed and burned")
-        .inner_html();
+        .inner_html()
 }

yew-functional/tests/use_context_hook.rs

@@ -0,0 +1,294 @@
+extern crate wasm_bindgen_test;
+
+use std::rc::Rc;
+use wasm_bindgen_test::*;
+wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
+
+extern crate yew;
+
+use yew::{html, App, Children, Html, Properties, Renderable};
+use yew_functional::{
+    use_context, use_effect, use_ref, use_state, ContextProvider, FunctionComponent,
+    FunctionProvider,
+};
+
+fn obtain_result(id: &str) -> String {
+    yew::utils::document()
+        .get_element_by_id(id)
+        .expect("No result found. Most likely, the application crashed and burned")
+        .inner_html()
+}
+
+#[wasm_bindgen_test]
+fn use_context_scoping_works() {
+    #[derive(Clone, Debug)]
+    struct ExampleContext(String);
+    struct UseContextFunctionOuter {}
+    struct UseContextFunctionInner {}
+    struct ExpectNoContextFunction {}
+    type UseContextComponent = FunctionComponent<UseContextFunctionOuter>;
+    type UseContextComponentInner = FunctionComponent<UseContextFunctionInner>;
+    type ExpectNoContextComponent = FunctionComponent<ExpectNoContextFunction>;
+    impl FunctionProvider for ExpectNoContextFunction {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            if use_context::<ExampleContext>().is_some() {
+                yew::services::ConsoleService::new().log(&format!(
+                    "Context should be None here, but was {:?}!",
+                    use_context::<ExampleContext>().unwrap()
+                ));
+            };
+            return html! {
+                <div></div>
+            };
+        }
+    }
+    impl FunctionProvider for UseContextFunctionOuter {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            type ExampleContextProvider = ContextProvider<ExampleContext>;
+            return html! {
+                <div>
+                    <ExampleContextProvider context=ExampleContext("wrong1".into())>
+                        <div>{"ignored"}</div>
+                    </ExampleContextProvider>
+                    <ExampleContextProvider context=ExampleContext("wrong2".into())>
+                        <ExampleContextProvider context=ExampleContext("correct".into())>
+                            <ExampleContextProvider context=ExampleContext("wrong1".into())>
+                                <div>{"ignored"}</div>
+                            </ExampleContextProvider>
+                            <UseContextComponentInner />
+                        </ExampleContextProvider>
+                    </ExampleContextProvider>
+                    <ExampleContextProvider context=ExampleContext("wrong3".into())>
+                        <div>{"ignored"}</div>
+                    </ExampleContextProvider>
+                    <ExpectNoContextComponent />
+                </div>
+            };
+        }
+    }
+    impl FunctionProvider for UseContextFunctionInner {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            let context = use_context::<ExampleContext>();
+            return html! {
+                <div id="result">{ &context.unwrap().0 }</div>
+            };
+        }
+    }
+
+    let app: App<UseContextComponent> = yew::App::new();
+    app.mount(yew::utils::document().get_element_by_id("output").unwrap());
+    let result: String = obtain_result("result");
+    assert_eq!("correct", result);
+}
+
+#[wasm_bindgen_test]
+fn use_context_works_with_multiple_types() {
+    #[derive(Clone, Debug, PartialEq)]
+    struct ContextA(u32);
+    #[derive(Clone, Debug, PartialEq)]
+    struct ContextB(u32);
+
+    struct Test1Function;
+    impl FunctionProvider for Test1Function {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            assert_eq!(use_context::<ContextA>(), Some(Rc::new(ContextA(2))));
+            assert_eq!(use_context::<ContextB>(), Some(Rc::new(ContextB(1))));
+
+            return html! {};
+        }
+    }
+    type Test1 = FunctionComponent<Test1Function>;
+
+    struct Test2Function;
+    impl FunctionProvider for Test2Function {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            assert_eq!(use_context::<ContextA>(), Some(Rc::new(ContextA(0))));
+            assert_eq!(use_context::<ContextB>(), Some(Rc::new(ContextB(1))));
+
+            return html! {};
+        }
+    }
+    type Test2 = FunctionComponent<Test2Function>;
+
+    struct Test3Function;
+    impl FunctionProvider for Test3Function {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            assert_eq!(use_context::<ContextA>(), Some(Rc::new(ContextA(0))));
+            assert_eq!(use_context::<ContextB>(), None);
+
+            return html! {};
+        }
+    }
+    type Test3 = FunctionComponent<Test3Function>;
+
+    struct Test4Function;
+    impl FunctionProvider for Test4Function {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            assert_eq!(use_context::<ContextA>(), None);
+            assert_eq!(use_context::<ContextB>(), None);
+
+            return html! {};
+        }
+    }
+    type Test4 = FunctionComponent<Test4Function>;
+
+    struct TestFunction;
+    impl FunctionProvider for TestFunction {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            type ContextAProvider = ContextProvider<ContextA>;
+            type ContextBProvider = ContextProvider<ContextB>;
+
+            return html! {
+                <div>
+                    <ContextAProvider context=ContextA(0)>
+                        <ContextBProvider context=ContextB(1)>
+                            <ContextAProvider context=ContextA(2)>
+                                <Test1/>
+                            </ContextAProvider>
+                            <Test2/>
+                        </ContextBProvider>
+                        <Test3/>
+                    </ContextAProvider>
+                    <Test4 />
+                </div>
+            };
+        }
+    }
+    type TestComponent = FunctionComponent<TestFunction>;
+
+    let app: App<TestComponent> = yew::App::new();
+    app.mount(yew::utils::document().get_element_by_id("output").unwrap());
+}
+
+#[wasm_bindgen_test]
+fn use_context_update_works() {
+    #[derive(Clone, Debug)]
+    struct MyContext(String);
+
+    #[derive(Clone, Debug, PartialEq, Properties)]
+    struct RenderCounterProps {
+        id: String,
+        children: Children,
+    }
+    struct RenderCounterFunction;
+    impl FunctionProvider for RenderCounterFunction {
+        type TProps = RenderCounterProps;
+
+        fn run(props: &Self::TProps) -> Html {
+            let counter = use_ref(|| 0);
+            *counter.borrow_mut() += 1;
+            return html! {
+                <>
+                    <div id=props.id.clone()>
+                        { format!("total: {}", counter.borrow()) }
+                    </div>
+                    { props.children.render() }
+                </>
+            };
+        }
+    }
+    type RenderCounter = FunctionComponent<RenderCounterFunction>;
+
+    #[derive(Clone, Debug, PartialEq, Properties)]
+    struct ContextOutletProps {
+        id: String,
+        #[prop_or_default]
+        magic: usize,
+    }
+    struct ContextOutletFunction;
+    impl FunctionProvider for ContextOutletFunction {
+        type TProps = ContextOutletProps;
+
+        fn run(props: &Self::TProps) -> Html {
+            let counter = use_ref(|| 0);
+            *counter.borrow_mut() += 1;
+
+            let ctx = use_context::<Rc<MyContext>>().expect("context not passed down");
+
+            return html! {
+                <>
+                    <div>{ format!("magic: {}\n", props.magic) }</div>
+                    <div id=props.id.clone()>
+                        { format!("current: {}, total: {}", ctx.0, counter.borrow()) }
+                    </div>
+                </>
+            };
+        }
+    }
+    type ContextOutlet = FunctionComponent<ContextOutletFunction>;
+
+    struct TestFunction;
+    impl FunctionProvider for TestFunction {
+        type TProps = ();
+
+        fn run(_props: &Self::TProps) -> Html {
+            type MyContextProvider = ContextProvider<Rc<MyContext>>;
+
+            let (ctx, set_ctx) = use_state(|| MyContext("hello".into()));
+            let rendered = use_ref(|| 0);
+
+            // this is used to force an update specific to test-2
+            let (magic_rc, set_magic) = use_state(|| 0);
+            let magic: usize = *magic_rc;
+
+            use_effect(move || {
+                let count = *rendered.borrow();
+                match count {
+                    0 => {
+                        set_ctx(MyContext("world".into()));
+                        *rendered.borrow_mut() += 1;
+                    }
+                    1 => {
+                        // force test-2 to re-render.
+                        set_magic(1);
+                        *rendered.borrow_mut() += 1;
+                    }
+                    2 => {
+                        set_ctx(MyContext("hello world!".into()));
+                        *rendered.borrow_mut() += 1;
+                    }
+                    _ => (),
+                };
+                || {}
+            });
+
+            return html! {
+                <MyContextProvider context=ctx.clone()>
+                    <RenderCounter id="test-0">
+                        <ContextOutlet id="test-1"/>
+                        <ContextOutlet id="test-2" magic=magic/>
+                    </RenderCounter>
+                </MyContextProvider>
+            };
+        }
+    }
+    type TestComponent = FunctionComponent<TestFunction>;
+
+    let app: App<TestComponent> = yew::App::new();
+    app.mount(yew::utils::document().get_element_by_id("output").unwrap());
+
+    // 1 initial render + 3 update steps
+    assert_eq!(obtain_result("test-0"), "total: 4");
+
+    // 1 initial + 2 context updates
+    assert_eq!(obtain_result("test-1"), "current: hello world!, total: 3");
+
+    // 1 initial + 1 context update + 1 magic update + 1 context update
+    assert_eq!(obtain_result("test-2"), "current: hello world!, total: 4");
+}