From cff277b3dbb0e0629c0461069c78475d7705ec0c Mon Sep 17 00:00:00 2001
From: Greg Johnston <greg.johnston@gmail.com>
Date: Tue, 19 Nov 2024 10:04:40 -0500
Subject: [PATCH] fix: add `Write` implementatations for `Field<T>` and
 `ArcField<T>` (closes #3257) (#3262)

---
 reactive_stores/src/arc_field.rs | 20 +++++++++++++++++++-
 reactive_stores/src/field.rs     | 30 ++++++++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/reactive_stores/src/arc_field.rs b/reactive_stores/src/arc_field.rs
index e21abc58cc..baa66eb4c8 100644
--- a/reactive_stores/src/arc_field.rs
+++ b/reactive_stores/src/arc_field.rs
@@ -7,6 +7,7 @@ use reactive_graph::{
     owner::Storage,
     traits::{
         DefinedAt, IsDisposed, Notify, ReadUntracked, Track, UntrackableGuard,
+        Write,
     },
 };
 use std::{
@@ -31,7 +32,8 @@ where
     trigger: StoreFieldTrigger,
     get_trigger: Arc<dyn Fn(StorePath) -> StoreFieldTrigger + Send + Sync>,
     read: Arc<dyn Fn() -> Option<StoreFieldReader<T>> + Send + Sync>,
-    write: Arc<dyn Fn() -> Option<StoreFieldWriter<T>> + Send + Sync>,
+    pub(crate) write:
+        Arc<dyn Fn() -> Option<StoreFieldWriter<T>> + Send + Sync>,
     keys: Arc<dyn Fn() -> Option<KeyMap> + Send + Sync>,
     track_field: Arc<dyn Fn() + Send + Sync>,
 }
@@ -329,6 +331,22 @@ impl<T> ReadUntracked for ArcField<T> {
     }
 }
 
+impl<T> Write for ArcField<T> {
+    type Value = T;
+
+    fn try_write(&self) -> Option<impl UntrackableGuard<Target = Self::Value>> {
+        (self.write)()
+    }
+
+    fn try_write_untracked(
+        &self,
+    ) -> Option<impl DerefMut<Target = Self::Value>> {
+        let mut guard = (self.write)()?;
+        guard.untrack();
+        Some(guard)
+    }
+}
+
 impl<T> IsDisposed for ArcField<T> {
     fn is_disposed(&self) -> bool {
         false
diff --git a/reactive_stores/src/field.rs b/reactive_stores/src/field.rs
index 89f0cc4678..d92a1e3fce 100644
--- a/reactive_stores/src/field.rs
+++ b/reactive_stores/src/field.rs
@@ -6,10 +6,18 @@ use crate::{
 };
 use reactive_graph::{
     owner::{ArenaItem, Storage, SyncStorage},
-    traits::{DefinedAt, IsDisposed, Notify, ReadUntracked, Track},
+    traits::{
+        DefinedAt, IsDisposed, Notify, ReadUntracked, Track, UntrackableGuard,
+        Write,
+    },
     unwrap_signal,
 };
-use std::{fmt::Debug, hash::Hash, ops::IndexMut, panic::Location};
+use std::{
+    fmt::Debug,
+    hash::Hash,
+    ops::{DerefMut, IndexMut},
+    panic::Location,
+};
 
 /// Wraps access to a single field of type `T`.
 ///
@@ -204,6 +212,24 @@ where
     }
 }
 
+impl<T> Write for Field<T> {
+    type Value = T;
+
+    fn try_write(&self) -> Option<impl UntrackableGuard<Target = Self::Value>> {
+        self.inner.try_get_value().and_then(|inner| (inner.write)())
+    }
+
+    fn try_write_untracked(
+        &self,
+    ) -> Option<impl DerefMut<Target = Self::Value>> {
+        self.inner.try_get_value().and_then(|inner| {
+            let mut guard = (inner.write)()?;
+            guard.untrack();
+            Some(guard)
+        })
+    }
+}
+
 impl<T, S> IsDisposed for Field<T, S> {
     fn is_disposed(&self) -> bool {
         self.inner.is_disposed()