From 5e57779435d3e8c8ce403ec7a8fcfb5739865e5d Mon Sep 17 00:00:00 2001
From: Evan Almloff <evanalmloff@gmail.com>
Date: Fri, 7 Jun 2024 02:01:42 +0200
Subject: [PATCH] fix memo and resource caller information (#2443)

---
 packages/hooks/src/use_resource.rs | 13 +++++++++++--
 packages/signals/src/memo.rs       | 23 +++++++++++++++--------
 packages/signals/src/signal.rs     | 10 ++++++++++
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/packages/hooks/src/use_resource.rs b/packages/hooks/src/use_resource.rs
index 4006dfb532..cce33ea5a8 100644
--- a/packages/hooks/src/use_resource.rs
+++ b/packages/hooks/src/use_resource.rs
@@ -68,15 +68,18 @@ use std::{cell::Cell, future::Future, rc::Rc};
 /// }
 /// ```
 #[must_use = "Consider using `cx.spawn` to run a future without reading its value"]
+#[track_caller]
 pub fn use_resource<T, F>(mut future: impl FnMut() -> F + 'static) -> Resource<T>
 where
     T: 'static,
     F: Future<Output = T> + 'static,
 {
+    let location = std::panic::Location::caller();
+
     let mut value = use_signal(|| None);
     let mut state = use_signal(|| UseResourceState::Pending);
     let (rc, changed) = use_hook(|| {
-        let (rc, changed) = ReactiveContext::new();
+        let (rc, changed) = ReactiveContext::new_with_origin(location);
         (rc, Rc::new(Cell::new(Some(changed))))
     });
 
@@ -92,7 +95,13 @@ where
 
             // Run each poll in the context of the reactive scope
             // This ensures the scope is properly subscribed to the future's dependencies
-            let res = future::poll_fn(|cx| rc.run_in(|| fut.poll_unpin(cx))).await;
+            let res = future::poll_fn(|cx| {
+                rc.run_in(|| {
+                    tracing::trace_span!("polling resource", location = %location)
+                        .in_scope(|| fut.poll_unpin(cx))
+                })
+            })
+            .await;
 
             // Set the value and state
             state.set(UseResourceState::Ready);
diff --git a/packages/signals/src/memo.rs b/packages/signals/src/memo.rs
index 7a1f047748..e5479c234b 100644
--- a/packages/signals/src/memo.rs
+++ b/packages/signals/src/memo.rs
@@ -5,7 +5,6 @@ use crate::{CopyValue, ReadOnlySignal};
 use std::{
     cell::RefCell,
     ops::Deref,
-    panic::Location,
     sync::{atomic::AtomicBool, Arc},
 };
 
@@ -36,7 +35,18 @@ where
 impl<T: 'static> Memo<T> {
     /// Create a new memo
     #[track_caller]
-    pub fn new(mut f: impl FnMut() -> T + 'static) -> Self
+    pub fn new(f: impl FnMut() -> T + 'static) -> Self
+    where
+        T: PartialEq,
+    {
+        Self::new_with_location(f, std::panic::Location::caller())
+    }
+
+    /// Create a new memo with an explicit location
+    pub fn new_with_location(
+        mut f: impl FnMut() -> T + 'static,
+        location: &'static std::panic::Location<'static>,
+    ) -> Self
     where
         T: PartialEq,
     {
@@ -50,11 +60,8 @@ impl<T: 'static> Memo<T> {
                 let _ = tx.unbounded_send(());
             }
         };
-        let rc = ReactiveContext::new_with_callback(
-            callback,
-            current_scope_id().unwrap(),
-            Location::caller(),
-        );
+        let rc =
+            ReactiveContext::new_with_callback(callback, current_scope_id().unwrap(), location);
 
         // Create a new signal in that context, wiring up its dependencies and subscribers
         let mut recompute = move || rc.run_in(&mut f);
@@ -64,7 +71,7 @@ impl<T: 'static> Memo<T> {
             dirty,
             callback: recompute,
         });
-        let state: Signal<T> = Signal::new(value);
+        let state: Signal<T> = Signal::new_with_caller(value, location);
 
         let memo = Memo {
             inner: state,
diff --git a/packages/signals/src/signal.rs b/packages/signals/src/signal.rs
index 28d11a852e..6bac531e0c 100644
--- a/packages/signals/src/signal.rs
+++ b/packages/signals/src/signal.rs
@@ -91,6 +91,16 @@ impl<T: PartialEq + 'static> Signal<T> {
     pub fn memo(f: impl FnMut() -> T + 'static) -> Memo<T> {
         Memo::new(f)
     }
+
+    /// Creates a new unsync Selector with an explicit location. The selector will be run immediately and whenever any signal it reads changes.
+    ///
+    /// Selectors can be used to efficiently compute derived data from signals.
+    pub fn memo_with_location(
+        f: impl FnMut() -> T + 'static,
+        location: &'static std::panic::Location<'static>,
+    ) -> Memo<T> {
+        Memo::new_with_location(f, location)
+    }
 }
 
 impl<T: 'static, S: Storage<SignalData<T>>> Signal<T, S> {