Deduplicate reactive scope updates/Reset subscriptions on reruns/fix use memo double update (#2506)

* deduplicate effect updates

* only subscribe to signals read in the current run of reactive scopes

* subscribe to memo reads after recomputing the value

packages/core/src/generational_box.rs

@@ -9,7 +9,7 @@ use std::{
 
 use generational_box::{AnyStorage, Owner, SyncStorage, UnsyncStorage};
 
-use crate::{innerlude::current_scope_id, ScopeId};
+use crate::{innerlude::current_scope_id, Runtime, ScopeId};
 
 /// Run a closure with the given owner.
 ///
@@ -91,12 +91,6 @@ pub fn current_owner<S: AnyStorage>() -> Owner<S> {
 impl ScopeId {
     /// Get the owner for the current scope.
     pub fn owner<S: AnyStorage>(self) -> Owner<S> {
-        match self.has_context() {
-            Some(rt) => rt,
-            None => {
-                let owner = S::owner();
-                self.provide_context(owner)
-            }
-        }
+        Runtime::with_scope(self, |cx| cx.owner::<S>()).expect("to be in a dioxus runtime")
     }
 }

packages/core/src/lib.rs

@@ -15,6 +15,7 @@ mod global_context;
 mod mutations;
 mod nodes;
 mod properties;
+mod reactive_context;
 mod render_signal;
 mod runtime;
 mod scheduler;
@@ -42,6 +43,7 @@ pub(crate) mod innerlude {
     pub use crate::mutations::*;
     pub use crate::nodes::*;
     pub use crate::properties::*;
+    pub use crate::reactive_context::*;
     pub use crate::runtime::{Runtime, RuntimeGuard};
     pub use crate::scheduler::*;
     pub use crate::scopes::*;
@@ -79,7 +81,7 @@ pub mod prelude {
         use_hook_with_cleanup, wait_for_next_render, with_owner, AnyValue, Attribute, Callback,
         Component, ComponentFunction, Element, ErrorBoundary, Event, EventHandler, Fragment,
         HasAttributes, IntoAttributeValue, IntoDynNode, OptionStringFromMarker, Properties,
-        Runtime, RuntimeGuard, ScopeId, ScopeState, SuperFrom, SuperInto, Task, Template,
-        TemplateAttribute, TemplateNode, Throw, VNode, VNodeInner, VirtualDom,
+        ReactiveContext, Runtime, RuntimeGuard, ScopeId, ScopeState, SuperFrom, SuperInto, Task,
+        Template, TemplateAttribute, TemplateNode, Throw, VNode, VNodeInner, VirtualDom,
     };
 }

packages/core/src/reactive_context.rs

@@ -0,0 +1,298 @@
+use crate::{
+    prelude::{current_scope_id, ScopeId},
+    scope_context::Scope,
+    tasks::SchedulerMsg,
+    Runtime,
+};
+use futures_channel::mpsc::UnboundedReceiver;
+use generational_box::{GenerationalBox, SyncStorage};
+use std::{
+    cell::RefCell,
+    collections::HashSet,
+    hash::Hash,
+    sync::{Arc, Mutex},
+};
+
+#[doc = include_str!("../docs/reactivity.md")]
+#[derive(Clone, Copy)]
+pub struct ReactiveContext {
+    scope: ScopeId,
+    inner: GenerationalBox<Inner, SyncStorage>,
+}
+
+impl PartialEq for ReactiveContext {
+    fn eq(&self, other: &Self) -> bool {
+        self.inner.ptr_eq(&other.inner)
+    }
+}
+
+impl Eq for ReactiveContext {}
+
+thread_local! {
+    static CURRENT: RefCell<Vec<ReactiveContext>> = const { RefCell::new(vec![]) };
+}
+
+impl std::fmt::Display for ReactiveContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        #[cfg(debug_assertions)]
+        {
+            if let Ok(read) = self.inner.try_read() {
+                if let Some(scope) = read.scope {
+                    return write!(f, "ReactiveContext(for scope: {:?})", scope);
+                }
+                return write!(f, "ReactiveContext created at {}", read.origin);
+            }
+        }
+        write!(f, "ReactiveContext")
+    }
+}
+
+impl ReactiveContext {
+    /// Create a new reactive context
+    #[track_caller]
+    pub fn new() -> (Self, UnboundedReceiver<()>) {
+        Self::new_with_origin(std::panic::Location::caller())
+    }
+
+    /// Create a new reactive context with a location for debugging purposes
+    /// This is useful for reactive contexts created within closures
+    pub fn new_with_origin(
+        origin: &'static std::panic::Location<'static>,
+    ) -> (Self, UnboundedReceiver<()>) {
+        let (tx, rx) = futures_channel::mpsc::unbounded();
+        let callback = move || {
+            // If there is already an update queued, we don't need to queue another
+            if !tx.is_empty() {
+                return;
+            }
+            let _ = tx.unbounded_send(());
+        };
+        let _self = Self::new_with_callback(callback, current_scope_id().unwrap(), origin);
+        (_self, rx)
+    }
+
+    /// Create a new reactive context that may update a scope. When any signal that this context subscribes to changes, the callback will be run
+    pub fn new_with_callback(
+        callback: impl FnMut() + Send + Sync + 'static,
+        scope: ScopeId,
+        #[allow(unused)] origin: &'static std::panic::Location<'static>,
+    ) -> Self {
+        let inner = Inner {
+            self_: None,
+            update: Box::new(callback),
+            subscribers: Default::default(),
+            #[cfg(debug_assertions)]
+            origin,
+            #[cfg(debug_assertions)]
+            scope: None,
+        };
+
+        let owner = scope.owner();
+
+        let self_ = Self {
+            scope,
+            inner: owner.insert(inner),
+        };
+
+        self_.inner.write().self_ = Some(self_);
+
+        self_
+    }
+
+    /// Get the current reactive context from the nearest reactive hook or scope
+    pub fn current() -> Option<Self> {
+        CURRENT.with(|current| current.borrow().last().cloned())
+    }
+
+    /// Create a reactive context for a scope id
+    pub(crate) fn new_for_scope(scope: &Scope, runtime: &Runtime) -> Self {
+        let id = scope.id;
+        let sender = runtime.sender.clone();
+        let update_scope = move || {
+            tracing::trace!("Marking scope {:?} as dirty", id);
+            sender.unbounded_send(SchedulerMsg::Immediate(id)).unwrap();
+        };
+
+        // Otherwise, create a new context at the current scope
+        let inner = Inner {
+            self_: None,
+            update: Box::new(update_scope),
+            subscribers: Default::default(),
+            #[cfg(debug_assertions)]
+            origin: std::panic::Location::caller(),
+            #[cfg(debug_assertions)]
+            scope: Some(id),
+        };
+
+        let owner = scope.owner();
+
+        let self_ = Self {
+            scope: id,
+            inner: owner.insert(inner),
+        };
+
+        self_.inner.write().self_ = Some(self_);
+
+        self_
+    }
+
+    /// Clear all subscribers to this context
+    pub fn clear_subscribers(&self) {
+        let old_subscribers = std::mem::take(&mut self.inner.write().subscribers);
+        for subscriber in old_subscribers {
+            subscriber.0.lock().unwrap().remove(self);
+        }
+    }
+
+    /// Update the subscribers
+    pub(crate) fn update_subscribers(&self) {
+        let subscribers = &self.inner.read().subscribers;
+        for subscriber in subscribers.iter() {
+            subscriber.0.lock().unwrap().insert(*self);
+        }
+    }
+
+    /// Reset the reactive context and then run the callback in the context. This can be used to create custom reactive hooks like `use_memo`.
+    ///
+    /// ```rust, no_run
+    /// # use dioxus::prelude::*;
+    /// # use futures_util::StreamExt;
+    /// fn use_simplified_memo(mut closure: impl FnMut() -> i32 + 'static) -> Signal<i32> {
+    ///     use_hook(|| {
+    ///         // Create a new reactive context and channel that will recieve a value every time a value the reactive context subscribes to changes
+    ///         let (reactive_context, mut changed) = ReactiveContext::new();
+    ///         // Compute the value of the memo inside the reactive context. This will subscribe the reactive context to any values you read inside the closure
+    ///         let value = reactive_context.reset_and_run_in(&mut closure);
+    ///         // Create a new signal with the value of the memo
+    ///         let mut signal = Signal::new(value);
+    ///         // Create a task that reruns the closure when the reactive context changes
+    ///         spawn(async move {
+    ///             while changed.next().await.is_some() {
+    ///                 // Since we reset the reactive context as we run the closure, our memo will only subscribe to the new values that are read in the closure
+    ///                 let new_value = reactive_context.run_in(&mut closure);
+    ///                 if new_value != value {
+    ///                     signal.set(new_value);
+    ///                 }
+    ///             }
+    ///         });
+    ///         signal
+    ///     })
+    /// }
+    ///
+    /// let mut boolean = use_signal(|| false);
+    /// let mut count = use_signal(|| 0);
+    /// // Because we use `reset_and_run_in` instead of just `run_in`, our memo will only subscribe to the signals that are read this run of the closure (initially just the boolean)
+    /// let memo = use_simplified_memo(move || if boolean() { count() } else { 0 });
+    /// println!("{memo}");
+    /// // Because the count signal is not read in this run of the closure, the memo will not rerun
+    /// count += 1;
+    /// println!("{memo}");
+    /// // Because the boolean signal is read in this run of the closure, the memo will rerun
+    /// boolean.toggle();
+    /// println!("{memo}");
+    /// // If we toggle the boolean again, and the memo unsubscribes from the count signal
+    /// boolean.toggle();
+    /// println!("{memo}");
+    /// ```
+    pub fn reset_and_run_in<O>(&self, f: impl FnOnce() -> O) -> O {
+        self.clear_subscribers();
+        self.run_in(f)
+    }
+
+    /// Run this function in the context of this reactive context
+    ///
+    /// This will set the current reactive context to this context for the duration of the function.
+    /// You can then get information about the current subscriptions.
+    pub fn run_in<O>(&self, f: impl FnOnce() -> O) -> O {
+        CURRENT.with(|current| current.borrow_mut().push(*self));
+        let out = f();
+        CURRENT.with(|current| current.borrow_mut().pop());
+        self.update_subscribers();
+        out
+    }
+
+    /// Marks this reactive context as dirty
+    ///
+    /// If there's a scope associated with this context, then it will be marked as dirty too
+    ///
+    /// Returns true if the context was marked as dirty, or false if the context has been dropped
+    pub fn mark_dirty(&self) -> bool {
+        if let Ok(mut self_write) = self.inner.try_write() {
+            #[cfg(debug_assertions)]
+            {
+                tracing::trace!(
+                    "Marking reactive context created at {} as dirty",
+                    self_write.origin
+                );
+            }
+
+            (self_write.update)();
+
+            true
+        } else {
+            false
+        }
+    }
+
+    /// Subscribe to this context. The reactive context will automatically remove itself from the subscriptions when it is reset.
+    pub fn subscribe(&self, subscriptions: Arc<Mutex<HashSet<ReactiveContext>>>) {
+        subscriptions.lock().unwrap().insert(*self);
+        self.inner
+            .write()
+            .subscribers
+            .insert(PointerHash(subscriptions));
+    }
+
+    /// Get the scope that inner CopyValue is associated with
+    pub fn origin_scope(&self) -> ScopeId {
+        self.scope
+    }
+}
+
+impl Hash for ReactiveContext {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.inner.id().hash(state);
+    }
+}
+
+struct PointerHash<T>(Arc<T>);
+
+impl<T> Hash for PointerHash<T> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        std::sync::Arc::<T>::as_ptr(&self.0).hash(state);
+    }
+}
+
+impl<T> PartialEq for PointerHash<T> {
+    fn eq(&self, other: &Self) -> bool {
+        std::sync::Arc::<T>::as_ptr(&self.0) == std::sync::Arc::<T>::as_ptr(&other.0)
+    }
+}
+
+impl<T> Eq for PointerHash<T> {}
+
+impl<T> Clone for PointerHash<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+type SubscriberMap = Mutex<HashSet<ReactiveContext>>;
+
+struct Inner {
+    self_: Option<ReactiveContext>,
+
+    // Futures will call .changed().await
+    update: Box<dyn FnMut() + Send + Sync>,
+
+    // Subscribers to this context
+    subscribers: HashSet<PointerHash<SubscriberMap>>,
+
+    // Debug information for signal subscriptions
+    #[cfg(debug_assertions)]
+    origin: &'static std::panic::Location<'static>,
+
+    #[cfg(debug_assertions)]
+    // The scope that this reactive context is associated with
+    scope: Option<ScopeId>,
+}

packages/core/src/scope_arena.rs

@@ -1,4 +1,5 @@
 use crate::innerlude::ScopeOrder;
+use crate::reactive_context::ReactiveContext;
 use crate::{
     any_props::{AnyProps, BoxedAnyProps},
     innerlude::ScopeState,
@@ -17,16 +18,18 @@ impl VirtualDom {
         let entry = self.scopes.vacant_entry();
         let id = ScopeId(entry.key());
 
+        let scope_runtime = Scope::new(name, id, parent_id, height);
+        let reactive_context = ReactiveContext::new_for_scope(&scope_runtime, &self.runtime);
+        self.runtime.create_scope(scope_runtime);
+
         let scope = entry.insert(ScopeState {
             runtime: self.runtime.clone(),
             context_id: id,
             props,
             last_rendered_node: Default::default(),
+            reactive_context,
         });
 
-        self.runtime
-            .create_scope(Scope::new(name, id, parent_id, height));
-
         scope
     }
 
@@ -52,7 +55,7 @@ impl VirtualDom {
             let props: &dyn AnyProps = &*scope.props;
 
             let span = tracing::trace_span!("render", scope = %scope.state().name);
-            span.in_scope(|| props.render())
+            span.in_scope(|| scope.reactive_context.reset_and_run_in(|| props.render()))
         };
 
         let context = scope.state();

packages/core/src/scope_context.rs

@@ -1,4 +1,5 @@
 use crate::{innerlude::SchedulerMsg, Element, Runtime, ScopeId, Task};
+use generational_box::{AnyStorage, Owner};
 use rustc_hash::FxHashSet;
 use std::{
     any::Any,
@@ -91,6 +92,17 @@ impl Scope {
         })
     }
 
+    /// Get the owner for the current scope.
+    pub fn owner<S: AnyStorage>(&self) -> Owner<S> {
+        match self.has_context() {
+            Some(rt) => rt,
+            None => {
+                let owner = S::owner();
+                self.provide_context(owner)
+            }
+        }
+    }
+
     /// Return any context of type T if it exists on this scope
     pub fn has_context<T: 'static + Clone>(&self) -> Option<T> {
         self.shared_contexts

packages/core/src/scopes.rs

@@ -1,5 +1,6 @@
 use crate::{
-    any_props::BoxedAnyProps, nodes::RenderReturn, runtime::Runtime, scope_context::Scope,
+    any_props::BoxedAnyProps, nodes::RenderReturn, reactive_context::ReactiveContext,
+    runtime::Runtime, scope_context::Scope,
 };
 use std::{cell::Ref, rc::Rc};
 
@@ -53,6 +54,7 @@ pub struct ScopeState {
     pub(crate) context_id: ScopeId,
     pub(crate) last_rendered_node: Option<RenderReturn>,
     pub(crate) props: BoxedAnyProps,
+    pub(crate) reactive_context: ReactiveContext,
 }
 
 impl Drop for ScopeState {