From 0a35b5a6705ee7528b47c81b1c6293ee768649c1 Mon Sep 17 00:00:00 2001
From: Evan Almloff <evanalmloff@gmail.com>
Date: Thu, 7 Dec 2023 13:21:47 -0600
Subject: [PATCH] fix render macro

---
 packages/core/src/arena.rs  | 11 ++++++++--
 packages/core/src/scopes.rs | 42 ++++++++++++++++++++++++++++++++++++-
 packages/rsx/src/lib.rs     |  4 +++-
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/packages/core/src/arena.rs b/packages/core/src/arena.rs
index ffdd4656b..ecefe20f7 100644
--- a/packages/core/src/arena.rs
+++ b/packages/core/src/arena.rs
@@ -1,4 +1,4 @@
-use std::ptr::NonNull;
+use crate::ScopeState;use std::ptr::NonNull;
 
 use crate::{
     innerlude::DirtyScope, nodes::RenderReturn, nodes::VNode, virtual_dom::VirtualDom,
@@ -174,7 +174,14 @@ impl VirtualDom {
         scope.borrowed_props.borrow_mut().clear();
 
         // Now that all the references are gone, we can safely drop our own references in our listeners.
-        let mut listeners = scope.attributes_to_drop_before_render.borrow_mut();
+        scope.drop_listeners();
+    }
+}
+
+impl ScopeState {
+    /// Drop all listeners that were allocated in the bump allocator.
+    pub(crate) fn drop_listeners(&self) {
+        let mut listeners = self.attributes_to_drop_before_render.borrow_mut();
         listeners.drain(..).for_each(|listener| {
             let listener = unsafe { &*listener };
             if let AttributeValue::Listener(l) = &listener.value {
diff --git a/packages/core/src/scopes.rs b/packages/core/src/scopes.rs
index 64a60b112..e3c42f515 100644
--- a/packages/core/src/scopes.rs
+++ b/packages/core/src/scopes.rs
@@ -102,6 +102,7 @@ pub struct ScopeState {
 
 impl Drop for ScopeState {
     fn drop(&mut self) {
+        self.drop_listeners();
         self.runtime.remove_context(self.context_id);
     }
 }
@@ -345,7 +346,7 @@ impl<'src> ScopeState {
     ///     // Actually build the tree and allocate it
     ///     cx.render(lazy_tree)
     /// }
-    ///```
+    /// ```
     pub fn render(&'src self, rsx: LazyNodes<'src, '_>) -> Element<'src> {
         let element = rsx.call(self);
 
@@ -387,6 +388,45 @@ impl<'src> ScopeState {
         Some(element)
     }
 
+    #[doc(hidden)]
+    /// Take a [`crate::VNode`] and link it to a [`ScopeState`]
+    /// 
+    /// This is used internally in the render macro
+    pub fn link_node(&'src self, element: &crate::VNode<'src>) {
+        let mut listeners = self.attributes_to_drop_before_render.borrow_mut();
+        for attr in element.dynamic_attrs {
+            match attr.value {
+                // We need to drop listeners before the next render because they may borrow data from the borrowed props which will be dropped
+                AttributeValue::Listener(_) => {
+                    let unbounded = unsafe { std::mem::transmute(attr as *const Attribute) };
+                    listeners.push(unbounded);
+                }
+                // We need to drop any values manually to make sure that their drop implementation is called before the next render
+                AttributeValue::Any(_) => {
+                    let unbounded = unsafe { std::mem::transmute(attr as *const Attribute) };
+                    self.previous_frame().add_attribute_to_drop(unbounded);
+                }
+
+                _ => (),
+            }
+        }
+
+        let mut props = self.borrowed_props.borrow_mut();
+        let mut drop_props = self
+            .previous_frame()
+            .props_to_drop_before_reset
+            .borrow_mut();
+        for node in element.dynamic_nodes {
+            if let DynamicNode::Component(comp) = node {
+                let unbounded = unsafe { std::mem::transmute(comp as *const VComponent) };
+                if !comp.static_props {
+                    props.push(unbounded);
+                }
+                drop_props.push(unbounded);
+            }
+        }
+    }
+
     /// Create a dynamic text node using [`Arguments`] and the [`ScopeState`]'s internal [`Bump`] allocator
     pub fn text_node(&'src self, args: Arguments) -> DynamicNode<'src> {
         DynamicNode::Text(VText {
diff --git a/packages/rsx/src/lib.rs b/packages/rsx/src/lib.rs
index aee1bb431..915d2bd2a 100644
--- a/packages/rsx/src/lib.rs
+++ b/packages/rsx/src/lib.rs
@@ -139,7 +139,9 @@ impl ToTokens for RenderCallBody {
         out_tokens.append_all(quote! {
             Some({
                 let __cx = cx;
-                #body
+                let __element = { #body };
+                cx.link_node(&__element);
+                __element
             })
         })
     }