From 07529e491cf85a5fd81d5c7718c43a41432c9468 Mon Sep 17 00:00:00 2001
From: Jonathan Kelley <jkelleyrtp@gmail.com>
Date: Sun, 24 Apr 2022 20:51:52 -0400
Subject: [PATCH] fix: subscription should apply internally to hook

---
 examples/fermi.rs                    |  3 ++-
 packages/fermi/src/hooks/atom_ref.rs | 40 +++++++++++++++++++---------
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/examples/fermi.rs b/examples/fermi.rs
index 9f87961b3..79e3da1d6 100644
--- a/examples/fermi.rs
+++ b/examples/fermi.rs
@@ -47,7 +47,8 @@ fn ChildWithRef(cx: Scope) -> Element {
                     cx.spawn(async move {
                         names.write().push("asd".to_string());
                     })
-                }
+                },
+                "Add name"
             }
         }
     })
diff --git a/packages/fermi/src/hooks/atom_ref.rs b/packages/fermi/src/hooks/atom_ref.rs
index f9c7f2869..fe488ee2f 100644
--- a/packages/fermi/src/hooks/atom_ref.rs
+++ b/packages/fermi/src/hooks/atom_ref.rs
@@ -16,15 +16,35 @@ use std::{
 pub fn use_atom_ref<T: 'static>(cx: &ScopeState, atom: AtomRef<T>) -> &UseAtomRef<T> {
     let root = use_atom_root(cx);
 
-    cx.use_hook(|_| {
+    &cx.use_hook(|_| {
         root.initialize(atom);
-        UseAtomRef {
-            ptr: atom.unique_id(),
-            root: root.clone(),
-            scope_id: cx.scope_id(),
-            value: root.register(atom, cx.scope_id()),
-        }
+        (
+            UseAtomRef {
+                ptr: atom.unique_id(),
+                root: root.clone(),
+                scope_id: cx.scope_id(),
+                value: root.register(atom, cx.scope_id()),
+            },
+            AtomRefSubscription {
+                ptr: atom.unique_id(),
+                root: root.clone(),
+                scope_id: cx.scope_id(),
+            },
+        )
     })
+    .0
+}
+
+pub struct AtomRefSubscription {
+    ptr: AtomId,
+    root: Rc<AtomRoot>,
+    scope_id: ScopeId,
+}
+
+impl Drop for AtomRefSubscription {
+    fn drop(&mut self) {
+        self.root.unsubscribe(self.ptr, self.scope_id)
+    }
 }
 
 pub struct UseAtomRef<T> {
@@ -65,9 +85,3 @@ impl<T: 'static> UseAtomRef<T> {
         self.root.set(self.ptr, new);
     }
 }
-
-impl<T> Drop for UseAtomRef<T> {
-    fn drop(&mut self) {
-        self.root.unsubscribe(self.ptr, self.scope_id)
-    }
-}