From be740b38eea09f08d32115e13582c2a59c7ff2aa Mon Sep 17 00:00:00 2001
From: Greg Johnston <greg.johnston@gmail.com>
Date: Sat, 30 Nov 2024 11:55:13 -0500
Subject: [PATCH] feat: add `NodeRef::on_load()` and writeable `NodeRef`
 (#3305)

---
 tachys/src/reactive_graph/node_ref.rs | 54 ++++++++++++++++++++++++++-
 1 file changed, 53 insertions(+), 1 deletion(-)

diff --git a/tachys/src/reactive_graph/node_ref.rs b/tachys/src/reactive_graph/node_ref.rs
index 30133f77b0..088bf2a63e 100644
--- a/tachys/src/reactive_graph/node_ref.rs
+++ b/tachys/src/reactive_graph/node_ref.rs
@@ -1,12 +1,17 @@
 use crate::html::{element::ElementType, node_ref::NodeRefContainer};
 use reactive_graph::{
+    effect::Effect,
     signal::{
         guards::{Derefable, ReadGuard},
         RwSignal,
     },
-    traits::{DefinedAt, ReadUntracked, Set, Track},
+    traits::{
+        DefinedAt, Get, Notify, ReadUntracked, Set, Track, UntrackableGuard,
+        Write,
+    },
 };
 use send_wrapper::SendWrapper;
+use std::{cell::Cell, ops::DerefMut};
 use wasm_bindgen::JsCast;
 
 /// A reactive reference to a DOM node that can be used with the `node_ref` attribute.
@@ -26,6 +31,25 @@ where
     pub fn new() -> Self {
         Self(RwSignal::new(None))
     }
+
+    /// Runs the provided closure when the `NodeRef` has been connected
+    /// with its element.
+    #[inline(always)]
+    pub fn on_load<F>(self, f: F)
+    where
+        E: 'static,
+        F: FnOnce(E::Output) + 'static,
+        E: ElementType,
+        E::Output: JsCast + Clone + 'static,
+    {
+        let f = Cell::new(Some(f));
+
+        Effect::new(move |_| {
+            if let Some(node_ref) = self.get() {
+                f.take().unwrap()(node_ref);
+            }
+        });
+    }
 }
 
 impl<E> Default for NodeRef<E>
@@ -78,6 +102,34 @@ where
     }
 }
 
+impl<E> Notify for NodeRef<E>
+where
+    E: ElementType,
+    E::Output: JsCast + Clone + 'static,
+{
+    fn notify(&self) {
+        self.0.notify();
+    }
+}
+
+impl<E> Write for NodeRef<E>
+where
+    E: ElementType,
+    E::Output: JsCast + Clone + 'static,
+{
+    type Value = Option<SendWrapper<E::Output>>;
+
+    fn try_write(&self) -> Option<impl UntrackableGuard<Target = Self::Value>> {
+        self.0.try_write()
+    }
+
+    fn try_write_untracked(
+        &self,
+    ) -> Option<impl DerefMut<Target = Self::Value>> {
+        self.0.try_write_untracked()
+    }
+}
+
 impl<E> ReadUntracked for NodeRef<E>
 where
     E: ElementType,