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,100 +27,131 @@ 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 ptr: *const _ = &val as &dyn FnOnce(NodeFactory<'a>) -> VNode<'a>;
+        let mut slot = Some(_val);
 
-            assert_eq!(
-                ptr as *const u8, &val as *const _ as *const u8,
-                "MISUSE: Closure returned different pointer"
-            );
-            assert_eq!(
-                std::mem::size_of_val(&*ptr),
-                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,
-                "BUG: Pointer layout is not (data_ptr, info...)"
-            );
+        let val = move |f| {
+            let inn = slot.take().unwrap();
+            inn(f)
+        };
 
-            // - Ensure that Self is aligned same as data requires
-            assert!(
-                std::mem::align_of::<F>() <= std::mem::align_of::<Self>(),
-                "TODO: Enforce alignment >{} (requires {})",
-                std::mem::align_of::<Self>(),
-                std::mem::align_of::<F>()
-            );
+        unsafe { LazyNodes::new_inner(val) }
+    }
 
-            let info = &words[1..];
-            let data = words[0] as *mut ();
-            let size = mem::size_of::<F>();
+    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,
+            "MISUSE: Closure returned different pointer"
+        );
+        assert_eq!(
+            std::mem::size_of_val(&*ptr),
+            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,
+            "BUG: Pointer layout is not (data_ptr, info...)"
+        );
+
+        // - Ensure that Self is aligned same as data requires
+        assert!(
+            std::mem::align_of::<F>() <= std::mem::align_of::<Self>(),
+            "TODO: Enforce alignment >{} (requires {})",
+            std::mem::align_of::<Self>(),
+            std::mem::align_of::<F>()
+        );
+
+        let info = &words[1..];
+        let data = words[0] as *mut ();
+        let size = mem::size_of::<F>();
+
+        let stored_size = info.len() * mem::size_of::<usize>() + size;
+        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! max: {}, actual: {:?}",
+                max_size,
+                stored_size
+            );
+            let mut buf: StackHeapSize = StackHeapSize::default();
 
-            if info.len() * mem::size_of::<usize>() + size > mem::size_of::<StackHeapSize>() {
-                log::debug!("lazy nodes was too large to fit into stack. falling back to heap");
+            assert!(info.len() + round_to_words(size) <= buf.as_ref().len());
 
-                Self {
-                    inner: StackNodeStorage::Heap(Box::new(val)),
-                }
-            } else {
-                log::debug!("lazy nodes fits on stack!");
-                let mut buf: StackHeapSize = StackHeapSize::default();
-
-                assert!(info.len() + round_to_words(size) <= buf.as_ref().len());
-
-                // Place pointer information at the end of the region
-                // - Allows the data to be at the start for alignment purposes
-                {
-                    let info_ofs = buf.as_ref().len() - info.len();
-                    let info_dst = &mut buf.as_mut()[info_ofs..];
-                    for (d, v) in Iterator::zip(info_dst.iter_mut(), info.iter()) {
-                        *d = *v;
-                    }
+            // Place pointer information at the end of the region
+            // - Allows the data to be at the start for alignment purposes
+            {
+                let info_ofs = buf.as_ref().len() - info.len();
+                let info_dst = &mut buf.as_mut()[info_ofs..];
+                for (d, v) in Iterator::zip(info_dst.iter_mut(), info.iter()) {
+                    *d = *v;
                 }
+            }
 
-                let src_ptr = data as *const u8;
-                let dataptr = buf.as_mut()[..].as_mut_ptr() as *mut u8;
-                for i in 0..size {
-                    *dataptr.add(i) = *src_ptr.add(i);
-                }
+            let src_ptr = data as *const u8;
+            let dataptr = buf.as_mut()[..].as_mut_ptr() as *mut u8;
+            for i in 0..size {
+                *dataptr.add(i) = *src_ptr.add(i);
+            }
 
-                std::mem::forget(val);
+            log::debug!("I am forgetting the contents of the stuff");
+            std::mem::forget(val);
 
-                Self {
-                    inner: StackNodeStorage::Stack(LazyStack { _align: [], buf }),
-                }
+            Self {
+                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::Stack(stack) => stack.call(f),
+            StackNodeStorage::Heap(mut lazy) => lazy(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..]);
-
-        Box::from_raw(g)
+        make_fat_ptr(data[..].as_ptr() as usize, &data[info_ofs..])
     }
 
-    fn call(mut self, f: NodeFactory) -> VNode {
-        let boxed = unsafe { self.create_boxed() };
-        boxed(f)
+    fn call<'a>(&mut self, f: NodeFactory<'a>) -> VNode<'a> {
+        let LazyStack { buf, .. } = self;
+        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() };
-        mem::drop(boxed);
+        if !self.dropped {
+            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>>
-    where
-        F: FnOnce(NodeFactory<'z>) -> VNode<'z> + 'b,
-    {
+    pub fn annotate_lazy<'z, 'b>(
+        f: impl FnOnce(NodeFactory<'z>) -> VNode<'z> + 'b,
+    ) -> Option<LazyNodes<'z, 'b>> {
         Some(LazyNodes::new(f))
     }
 }