From 4a83ffca6f3240db07e3b437092f65953d7034dc Mon Sep 17 00:00:00 2001
From: Greg Johnston <greg.johnston@gmail.com>
Date: Thu, 19 Oct 2023 15:56:57 -0400
Subject: [PATCH] fix: `try_update()` and `try_set()` on `Resource` should not
 panic (closes #1915) (#1916)

---
 leptos_reactive/src/resource.rs |  8 +++++---
 leptos_reactive/src/runtime.rs  | 32 +++++++++++++++++++++++++-------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/leptos_reactive/src/resource.rs b/leptos_reactive/src/resource.rs
index 329f0899ba..43c7f39d7c 100644
--- a/leptos_reactive/src/resource.rs
+++ b/leptos_reactive/src/resource.rs
@@ -685,7 +685,7 @@ impl<S, T> SignalUpdate for Resource<S, T> {
     #[inline(always)]
     fn try_update<O>(&self, f: impl FnOnce(&mut Option<T>) -> O) -> Option<O> {
         with_runtime(|runtime| {
-            runtime.resource(self.id, |resource: &ResourceState<S, T>| {
+            runtime.try_resource(self.id, |resource: &ResourceState<S, T>| {
                 if resource.loading.get_untracked() {
                     resource.version.set(resource.version.get() + 1);
                     for suspense_context in
@@ -703,6 +703,7 @@ impl<S, T> SignalUpdate for Resource<S, T> {
         })
         .ok()
         .flatten()
+        .flatten()
     }
 }
 
@@ -860,8 +861,9 @@ impl<S, T> SignalSet for Resource<S, T> {
     )]
     #[inline(always)]
     fn try_set(&self, new_value: T) -> Option<T> {
-        self.update(|n| *n = Some(new_value));
-        None
+        let mut new_value = Some(new_value);
+        self.try_update(|n| *n = new_value.take());
+        new_value
     }
 }
 
diff --git a/leptos_reactive/src/runtime.rs b/leptos_reactive/src/runtime.rs
index e9a0275c8f..88bca3c2ec 100644
--- a/leptos_reactive/src/runtime.rs
+++ b/leptos_reactive/src/runtime.rs
@@ -1249,6 +1249,7 @@ impl Runtime {
             .borrow_mut()
             .insert(AnyResource::Serializable(state))
     }
+
     #[cfg_attr(
         any(debug_assertions, feature = "ssr"),
         instrument(level = "trace", skip_all,)
@@ -1259,13 +1260,35 @@ impl Runtime {
         id: ResourceId,
         f: impl FnOnce(&ResourceState<S, T>) -> U,
     ) -> U
+    where
+        S: 'static,
+        T: 'static,
+    {
+        self.try_resource(id, f).unwrap_or_else(|| {
+            panic!(
+                "couldn't locate {id:?} at {:?}",
+                std::panic::Location::caller()
+            );
+        })
+    }
+
+    #[cfg_attr(
+        any(debug_assertions, feature = "ssr"),
+        instrument(level = "trace", skip_all,)
+    )]
+    #[track_caller]
+    pub(crate) fn try_resource<S, T, U>(
+        &self,
+        id: ResourceId,
+        f: impl FnOnce(&ResourceState<S, T>) -> U,
+    ) -> Option<U>
     where
         S: 'static,
         T: 'static,
     {
         let resources = { self.resources.borrow().clone() };
         let res = resources.get(id);
-        if let Some(res) = res {
+        res.map(|res| {
             let res_state = match res {
                 AnyResource::Unserializable(res) => res.as_any(),
                 AnyResource::Serializable(res) => res.as_any(),
@@ -1281,12 +1304,7 @@ impl Runtime {
                     std::any::type_name::<T>(),
                 );
             }
-        } else {
-            panic!(
-                "couldn't locate {id:?} at {:?}",
-                std::panic::Location::caller()
-            );
-        }
+        })
     }
 
     /// Returns IDs for all [`Resource`](crate::Resource)s found on any scope.