Bring context to standard components

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

@@ -1,174 +1,66 @@
 use crate::{get_current_scope, use_hook};
-use std::any::TypeId;
-use std::cell::RefCell;
-use std::rc::{Rc, Weak};
-use std::{iter, mem};
-use yew::html;
-use yew::html::{AnyScope, Scope};
-use yew::{Children, Component, ComponentLink, Html, Properties};
-
-type ConsumerCallback<T> = Box<dyn Fn(Rc<T>)>;
-type UseContextOutput<T> = Option<Rc<T>>;
-
-struct UseContext<T2: Clone + PartialEq + 'static> {
-    provider_scope: Option<Scope<ContextProvider<T2>>>,
-    current_context: Option<Rc<T2>>,
-    callback: Option<Rc<ConsumerCallback<T2>>>,
-}
+use yew::context::ContextHandle;
 
 /// Hook for consuming context values in function components.
-pub fn use_context<T: Clone + PartialEq + 'static>() -> UseContextOutput<T> {
+/// The context of the type passed as `T` is returned. If there is no such context in scope, `None` is returned.
+/// A component which calls `use_context` will re-render when the data of the context changes.
+///
+/// More information about contexts and how to define and consume them can be found on [Yew Docs](https://yew.rs).
+///
+/// # Example
+/// ```rust
+/// # use yew_functional::{function_component, use_context};
+/// # use yew::prelude::*;
+/// # use std::rc::Rc;
+///
+/// # #[derive(Clone, Debug, PartialEq)]
+/// # struct ThemeContext {
+/// #    foreground: String,
+/// #    background: String,
+/// # }
+/// #[function_component(ThemedButton)]
+/// pub fn themed_button() -> Html {
+///     let theme = use_context::<Rc<ThemeContext>>().expect("no ctx found");
+///
+///     html! {
+///         <button style=format!("background: {}; color: {}", theme.background, theme.foreground)>
+///             { "Click me" }
+///         </button>
+///     }
+/// }
+/// ```
+pub fn use_context<T: Clone + PartialEq + 'static>() -> Option<T> {
+    struct UseContextState<T2: Clone + PartialEq + 'static> {
+        initialized: bool,
+        context: Option<(T2, ContextHandle<T2>)>,
+    }
+
     let scope = get_current_scope()
         .expect("No current Scope. `use_context` can only be called inside function components");
 
     use_hook(
-        // Initializer
-        move || {
-            let provider_scope = find_context_provider_scope::<T>(&scope);
-            let current_context =
-                with_provider_component(&provider_scope, |comp| Rc::clone(&comp.context));
-
-            UseContext {
-                provider_scope,
-                current_context,
-                callback: None,
-            }
+        move || UseContextState {
+            initialized: false,
+            context: None,
         },
-        // Runner
-        |hook, updater| {
-            // setup a listener for the context provider to update us
-            let listener = move |ctx: Rc<T>| {
-                updater.callback(move |state: &mut UseContext<T>| {
-                    state.current_context = Some(ctx);
-                    true
-                });
-            };
-            hook.callback = Some(Rc::new(Box::new(listener)));
-
-            // Subscribe to the context provider with our callback
-            let weak_cb = Rc::downgrade(hook.callback.as_ref().unwrap());
-            with_provider_component(&hook.provider_scope, |comp| {
-                comp.subscribe_consumer(weak_cb)
-            });
+        |state: &mut UseContextState<T>, updater| {
+            if !state.initialized {
+                state.initialized = true;
+                let callback = move |ctx: T| {
+                    updater.callback(|state: &mut UseContextState<T>| {
+                        if let Some(context) = &mut state.context {
+                            context.0 = ctx;
+                        }
+                        true
+                    });
+                };
+                state.context = scope.context::<T>(callback.into());
+            }
 
-            // Return the current state
-            hook.current_context.clone()
+            Some(state.context.as_ref()?.0.clone())
         },
-        // Cleanup
-        |hook| {
-            if let Some(cb) = hook.callback.take() {
-                drop(cb);
-            }
+        |state| {
+            state.context = None;
         },
     )
 }
-
-/// Props for [`ContextProvider`]
-#[derive(Clone, PartialEq, Properties)]
-pub struct ContextProviderProps<T: Clone + PartialEq> {
-    pub context: T,
-    pub children: Children,
-}
-
-/// The context provider component.
-///
-/// Every child (direct or indirect) of this component may access the context value.
-/// Currently the only way to consume the context is using the [`use_context`] hook.
-pub struct ContextProvider<T: Clone + PartialEq + 'static> {
-    context: Rc<T>,
-    children: Children,
-    consumers: RefCell<Vec<Weak<ConsumerCallback<T>>>>,
-}
-
-impl<T: Clone + PartialEq> 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>>) {
-        // consumers re-subscribe on every render. Try to keep the subscriber list small by reusing dead slots.
-        let mut consumers = self.consumers.borrow_mut();
-        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 + PartialEq + 'static> Component for ContextProvider<T> {
-    type Message = ();
-    type Properties = ContextProviderProps<T>;
-
-    fn create(props: Self::Properties, _link: ComponentLink<Self>) -> Self {
-        Self {
-            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 {
-            false
-        } else {
-            self.children = props.children;
-            true
-        };
-
-        let new_context = Rc::new(props.context);
-        if self.context != new_context {
-            self.context = new_context;
-            self.notify_consumers();
-        }
-
-        should_render
-    }
-
-    fn view(&self) -> Html {
-        html! { <>{ self.children.clone() }</> }
-    }
-}
-
-fn find_context_provider_scope<T: Clone + PartialEq + 'static>(
-    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 + PartialEq,
-    F: FnOnce(&ContextProvider<T>) -> R,
-{
-    provider_scope
-        .as_ref()
-        .and_then(|scope| scope.get_component().map(|comp| f(&*comp)))
-}

packages/yew-functional/tests/use_context.rs

@@ -3,10 +3,9 @@ mod common;
 use common::obtain_result_by_id;
 use std::rc::Rc;
 use wasm_bindgen_test::*;
-use yew::{html, App, Children, Html, Properties};
+use yew::{html, App, Children, ContextProvider, Html, Properties};
 use yew_functional::{
-    use_context, use_effect, use_ref, use_state, ContextProvider, FunctionComponent,
-    FunctionProvider,
+    use_context, use_effect, use_ref, use_state, FunctionComponent, FunctionProvider,
 };
 
 wasm_bindgen_test::wasm_bindgen_test_configure!(run_in_browser);
@@ -91,8 +90,8 @@ fn use_context_works_with_multiple_types() {
         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))));
+            assert_eq!(use_context::<ContextA>(), Some(ContextA(2)));
+            assert_eq!(use_context::<ContextB>(), Some(ContextB(1)));
 
             return html! {};
         }
@@ -104,8 +103,8 @@ fn use_context_works_with_multiple_types() {
         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))));
+            assert_eq!(use_context::<ContextA>(), Some(ContextA(0)));
+            assert_eq!(use_context::<ContextB>(), Some(ContextB(1)));
 
             return html! {};
         }
@@ -117,7 +116,7 @@ fn use_context_works_with_multiple_types() {
         type TProps = ();
 
         fn run(_props: &Self::TProps) -> Html {
-            assert_eq!(use_context::<ContextA>(), Some(Rc::new(ContextA(0))));
+            assert_eq!(use_context::<ContextA>(), Some(ContextA(0)));
             assert_eq!(use_context::<ContextB>(), None);
 
             return html! {};

packages/yew/src/context.rs

@@ -0,0 +1,111 @@
+//! This module defines the `ContextProvider` component.
+
+use crate::html::Scope;
+use crate::{html, Callback, Children, Component, ComponentLink, Html, Properties};
+use slab::Slab;
+use std::cell::RefCell;
+
+/// Props for [`ContextProvider`]
+#[derive(Debug, Clone, PartialEq, Properties)]
+pub struct ContextProviderProps<T: Clone + PartialEq> {
+    /// Context value to be passed down
+    pub context: T,
+    /// Children
+    pub children: Children,
+}
+
+/// The context provider component.
+///
+/// Every child (direct or indirect) of this component may access the context value.
+/// Currently the only way to consume the context is using the [`use_context`] hook.
+#[derive(Debug)]
+pub struct ContextProvider<T: Clone + PartialEq + 'static> {
+    link: ComponentLink<Self>,
+    context: T,
+    children: Children,
+    consumers: RefCell<Slab<Callback<T>>>,
+}
+
+/// Owns the connection to a context provider. When dropped, the component will
+/// no longer receive updates from the provider.
+#[derive(Debug)]
+pub struct ContextHandle<T: Clone + PartialEq + 'static> {
+    provider: Scope<ContextProvider<T>>,
+    key: usize,
+}
+
+impl<T: Clone + PartialEq + 'static> Drop for ContextHandle<T> {
+    fn drop(&mut self) {
+        if let Some(component) = self.provider.get_component() {
+            component.consumers.borrow_mut().remove(self.key);
+        }
+    }
+}
+
+impl<T: Clone + PartialEq> 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.
+    pub(crate) fn subscribe_consumer(&self, callback: Callback<T>) -> (T, ContextHandle<T>) {
+        let ctx = self.context.clone();
+        let key = self.consumers.borrow_mut().insert(callback);
+
+        (
+            ctx,
+            ContextHandle {
+                provider: self.link.clone(),
+                key,
+            },
+        )
+    }
+
+    /// Notify all subscribed consumers and remove dropped consumers from the list.
+    fn notify_consumers(&mut self) {
+        let consumers: Vec<Callback<T>> = self
+            .consumers
+            .borrow()
+            .iter()
+            .map(|(_, v)| v.clone())
+            .collect();
+        for consumer in consumers {
+            consumer.emit(self.context.clone());
+        }
+    }
+}
+
+impl<T: Clone + PartialEq + 'static> Component for ContextProvider<T> {
+    type Message = ();
+    type Properties = ContextProviderProps<T>;
+
+    fn create(props: Self::Properties, link: ComponentLink<Self>) -> Self {
+        Self {
+            link,
+            children: props.children,
+            context: props.context,
+            consumers: RefCell::new(Slab::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 {
+            false
+        } else {
+            self.children = props.children;
+            true
+        };
+
+        if self.context != props.context {
+            self.context = props.context;
+            self.notify_consumers();
+        }
+
+        should_render
+    }
+
+    fn view(&self) -> Html {
+        html! { <>{ self.children.clone() }</> }
+    }
+}