wip: dst nodes leak but work

packages/core/src/lazynodes.rs

@@ -12,7 +12,7 @@
 //! support non-static closures, so we've implemented the core logic of `ValueA` in this module.
 
 use crate::prelude::{NodeFactory, VNode};
-use std::mem;
+use std::mem::{self, ManuallyDrop};
 
 /// A concrete type provider for closures that build VNode structures.
 ///
@@ -27,20 +27,33 @@ pub struct LazyNodes<'a, 'b> {
     inner: StackNodeStorage<'a, 'b>,
 }
 
-type StackHeapSize = [usize; 0];
+type StackHeapSize = [usize; 8];
 
 enum StackNodeStorage<'a, 'b> {
     Stack(LazyStack),
-    Heap(Box<dyn FnOnce(NodeFactory<'a>) -> VNode<'a> + 'b>),
+    Heap(Box<dyn FnMut(NodeFactory<'a>) -> VNode<'a> + 'b>),
 }
 
 impl<'a, 'b> LazyNodes<'a, 'b> {
-    pub fn new<F>(val: F) -> Self
+    pub fn new<F>(_val: F) -> Self
     where
         F: FnOnce(NodeFactory<'a>) -> VNode<'a> + 'b,
     {
-        unsafe {
+        let mut slot = Some(_val);
-            let mut ptr: *const _ = &val as &dyn FnOnce(NodeFactory<'a>) -> VNode<'a>;
+
+        let val = move |f| {
+            let inn = slot.take().unwrap();
+            inn(f)
+        };
+
+        unsafe { LazyNodes::new_inner(val) }
+    }
+
+    unsafe fn new_inner<F>(val: F) -> Self
+    where
+        F: FnMut(NodeFactory<'a>) -> VNode<'a> + 'b,
+    {
+        let mut ptr: *const _ = &val as &dyn FnMut(NodeFactory<'a>) -> VNode<'a>;
 
         assert_eq!(
             ptr as *const u8, &val as *const _ as *const u8,
@@ -51,6 +64,7 @@ impl<'a, 'b> LazyNodes<'a, 'b> {
             std::mem::size_of::<F>(),
             "MISUSE: Closure returned a subset pointer"
         );
+
         let words = ptr_as_slice(&mut ptr);
         assert!(
             words[0] == &val as *const _ as usize,
@@ -69,14 +83,25 @@ impl<'a, 'b> LazyNodes<'a, 'b> {
         let data = words[0] as *mut ();
         let size = mem::size_of::<F>();
 
-            if info.len() * mem::size_of::<usize>() + size > mem::size_of::<StackHeapSize>() {
+        let stored_size = info.len() * mem::size_of::<usize>() + size;
-                log::debug!("lazy nodes was too large to fit into stack. falling back to heap");
+        let max_size = mem::size_of::<StackHeapSize>();
+
+        if stored_size > max_size {
+            log::debug!(
+                    "lazy nodes was too large to fit into stack. falling back to heap. max: {}, actual {:?}",
+                    max_size,
+                    stored_size
+                );
 
             Self {
                 inner: StackNodeStorage::Heap(Box::new(val)),
             }
         } else {
-                log::debug!("lazy nodes fits on stack!");
+            log::debug!(
+                "lazy nodes fits on stack! max: {}, actual: {:?}",
+                max_size,
+                stored_size
+            );
             let mut buf: StackHeapSize = StackHeapSize::default();
 
             assert!(info.len() + round_to_words(size) <= buf.as_ref().len());
@@ -97,30 +122,36 @@ impl<'a, 'b> LazyNodes<'a, 'b> {
                 *dataptr.add(i) = *src_ptr.add(i);
             }
 
+            log::debug!("I am forgetting the contents of the stuff");
             std::mem::forget(val);
 
             Self {
-                    inner: StackNodeStorage::Stack(LazyStack { _align: [], buf }),
+                inner: StackNodeStorage::Stack(LazyStack {
-                }
+                    _align: [],
+                    buf,
+                    dropped: false,
+                }),
             }
         }
     }
 
-    pub fn call(self, f: NodeFactory<'a>) -> VNode<'a> {
+    pub fn call(mut self, f: NodeFactory<'a>) -> VNode<'a> {
         match self.inner {
-            StackNodeStorage::Heap(lazy) => lazy(f),
+            StackNodeStorage::Heap(mut lazy) => lazy(f),
-            StackNodeStorage::Stack(stack) => stack.call(f),
+            StackNodeStorage::Stack(mut stack) => stack.call(f),
         }
+        // todo drop?
     }
 }
 
 struct LazyStack {
     _align: [u64; 0],
     buf: StackHeapSize,
+    dropped: bool,
 }
 
 impl LazyStack {
-    unsafe fn create_boxed<'a>(&mut self) -> Box<dyn FnOnce(NodeFactory<'a>) -> VNode<'a>> {
+    unsafe fn as_raw<'a>(&mut self) -> *mut dyn FnOnce(NodeFactory<'a>) -> VNode<'a> {
         let LazyStack { buf, .. } = self;
         let data = buf.as_ref();
 
@@ -130,21 +161,44 @@ impl LazyStack {
 
         let info_ofs = data.len() - info_size;
 
-        let g: *mut dyn FnOnce(NodeFactory<'a>) -> VNode<'a> =
+        make_fat_ptr(data[..].as_ptr() as usize, &data[info_ofs..])
-            make_fat_ptr(data[..].as_ptr() as usize, &data[info_ofs..]);
-
-        Box::from_raw(g)
     }
 
-    fn call(mut self, f: NodeFactory) -> VNode {
+    fn call<'a>(&mut self, f: NodeFactory<'a>) -> VNode<'a> {
-        let boxed = unsafe { self.create_boxed() };
+        let LazyStack { buf, .. } = self;
-        boxed(f)
+        let data = buf.as_ref();
+
+        let info_size = mem::size_of::<*mut dyn FnOnce(NodeFactory<'a>) -> VNode<'a>>()
+            / mem::size_of::<usize>()
+            - 1;
+
+        let info_ofs = data.len() - info_size;
+
+        let g: *mut dyn FnMut(NodeFactory<'a>) -> VNode<'a> =
+            unsafe { make_fat_ptr(data[..].as_ptr() as usize, &data[info_ofs..]) };
+
+        self.dropped = true;
+
+        let clos = unsafe { &mut *g };
+        clos(f)
     }
 }
 impl Drop for LazyStack {
     fn drop(&mut self) {
-        let boxed = unsafe { self.create_boxed() };
+        if !self.dropped {
-        mem::drop(boxed);
+            log::debug!("manually dropping lazy nodes");
+
+            // match &mut self.inner {
+            //     StackNodeStorage::Heap(mut lazy) => {
+            //         log::debug!("dropping lazy nodes on heap");
+            //         std::mem::drop(lazy);
+            //     }
+            //     StackNodeStorage::Stack(mut stack) => {
+            //         log::debug!("dropping lazy nodes on stack");
+            //         std::mem::drop(stack);
+            //     }
+            // }
+        }
     }
 }
 
@@ -191,3 +245,39 @@ fn it_works() {
 
     dbg!(g);
 }
+
+#[test]
+fn it_drops() {
+    let bump = bumpalo::Bump::new();
+
+    simple_logger::init().unwrap();
+
+    let factory = NodeFactory { bump: &bump };
+
+    struct DropInner {
+        id: i32,
+    }
+    impl Drop for DropInner {
+        fn drop(&mut self) {
+            log::debug!("dropping inner");
+        }
+    }
+
+    let it = (0..10).map(|i| {
+        NodeFactory::annotate_lazy(move |f| {
+            log::debug!("hell closure");
+            let inner = DropInner { id: i };
+            f.text(format_args!("hello world {:?}", inner.id))
+        })
+    });
+
+    let caller = NodeFactory::annotate_lazy(|f| {
+        log::debug!("main closure");
+        f.fragment_from_iter(it)
+    })
+    .unwrap();
+
+    let nodes = caller.call(factory);
+
+    dbg!(nodes);
+}

packages/core/src/nodes.rs

@@ -679,10 +679,9 @@ impl<'a> NodeFactory<'a> {
         })
     }
 
-    pub fn annotate_lazy<'z, 'b, F>(f: F) -> Option<LazyNodes<'z, 'b>>
+    pub fn annotate_lazy<'z, 'b>(
-    where
+        f: impl FnOnce(NodeFactory<'z>) -> VNode<'z> + 'b,
-        F: FnOnce(NodeFactory<'z>) -> VNode<'z> + 'b,
+    ) -> Option<LazyNodes<'z, 'b>> {
-    {
         Some(LazyNodes::new(f))
     }
 }