Remove nodes on handle drop

matthunz · Dec 11, 2023 · 1c112a1 · 1c112a1
1 parent 50b1300
commit 1c112a1
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 17 deletions.
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
@@ -74,7 +74,7 @@ pub fn object(_attrs: TokenStream, input: TokenStream) -> TokenStream {
                 let ident = sig.ident;
                 sender_items.push(parse_quote! {
                     pub fn #ident(&self) -> viewbuilder::Signal<(#(#input_tys),*,)> {
-                        viewbuilder::Signal::new(self.handle.key)
+                        viewbuilder::Signal::new(self.handle.key())
                     }
                 });
             }

diff --git a/src/handle.rs b/src/handle.rs
@@ -23,6 +23,10 @@ impl<O: Object> Clone for Handle<O> {
 }
 
 impl<O: Object> Handle<O> {
+    pub fn key(&self) -> DefaultKey {
+        self.state.key()
+    }
+
     /// Send an update to the object.
     pub fn update(&self, f: impl FnMut(&mut O) + 'static)
     where
@@ -45,35 +49,56 @@ impl<O: Object> Deref for Handle<O> {
     }
 }
 
+struct Dropper {
+    key: DefaultKey,
+}
+
+impl Drop for Dropper {
+    fn drop(&mut self) {
+        Runtime::current().inner.borrow_mut().nodes.remove(self.key);
+    }
+}
+
 pub struct HandleState<O: Object> {
-    pub key: DefaultKey,
-    pub(crate) _marker: PhantomData<O>,
+    dropper: Rc<Dropper>,
+    _marker: PhantomData<O>,
 }
 
 impl<O: Object> Clone for HandleState<O> {
     fn clone(&self) -> Self {
         Self {
-            key: self.key.clone(),
+            dropper: self.dropper.clone(),
             _marker: self._marker.clone(),
         }
     }
 }
 
 impl<O: Object> HandleState<O> {
+    pub(crate) fn new(key: DefaultKey) -> Self {
+        Self {
+            dropper: Rc::new(Dropper { key }),
+            _marker: PhantomData,
+        }
+    }
+
+    pub fn key(&self) -> DefaultKey {
+        self.dropper.key
+    }
+
     /// Send an update to the object.
     pub fn update(&self, mut f: impl FnMut(&mut O) + 'static)
     where
         O: 'static,
     {
         Runtime::current().inner.borrow_mut().updates.push((
-            self.key,
+            self.key(),
             Box::new(move |element| f(element.downcast_mut().unwrap())),
         ))
     }
 
     /// Immutably borrow the object.
     pub fn borrow(&self) -> Ref<O> {
-        let guard = Runtime::current().inner.borrow().nodes[self.key]
+        let guard = Runtime::current().inner.borrow().nodes[self.key()]
             .object
             .clone();
 

diff --git a/src/lib.rs b/src/lib.rs
@@ -48,7 +48,8 @@ impl App {
 
     pub fn insert_window(&mut self, handle: Handle<Window>) {
         let window = winit::window::Window::new(&self.event_loop).unwrap();
-        self.windows.insert(window.id(), (window, handle.state.key));
+        self.windows
+            .insert(window.id(), (window, handle.state.key()));
     }
 
     pub fn run(self) {

diff --git a/src/object.rs b/src/object.rs
@@ -1,5 +1,5 @@
 use crate::{Handle, HandleState, Node, Runtime};
-use std::{cell::RefCell, marker::PhantomData, rc::Rc};
+use std::{cell::RefCell, rc::Rc};
 
 /// A reactive object.
 pub trait Object: Sized {
@@ -17,15 +17,8 @@ pub trait Object: Sized {
         });
 
         Handle {
-            state: HandleState {
-                key,
-                _marker: PhantomData,
-            },
-            sender: HandleState {
-                key,
-                _marker: PhantomData,
-            }
-            .into(),
+            state: HandleState::new(key),
+            sender: HandleState::new(key).into(),
         }
     }
 }