wip: implement dst magic

packages/core/src/lazynodes.rs

@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use std::mem;
 
 /*
 Remember: calls to rsx! are lazy - they are not evaluated immediately.
@@ -15,7 +15,7 @@ Our solution is to try and manually allocate the closure onto the stack.
 If it fails, then we default to Box.
 
 */
-use crate::innerlude::{IntoVNode, NodeFactory, VNode};
+use crate::innerlude::{NodeFactory, VNode};
 
 /// A concrete type provider for closures that build VNode structures.
 ///
@@ -30,103 +30,146 @@ pub struct LazyNodes<'a, 'b> {
     inner: StackNodeStorage<'a, 'b>,
 }
 
+type StackHeapSize = [usize; 12];
+
+enum StackNodeStorage<'a, 'b> {
+    Stack(LazyStack),
+    Heap(Box<dyn FnOnce(NodeFactory<'a>) -> VNode<'a> + 'b>),
+}
+
 impl<'a, 'b> LazyNodes<'a, 'b> {
-    pub fn new<F>(f: F) -> Self
+    pub fn new<F>(val: F) -> Self
     where
         F: FnOnce(NodeFactory<'a>) -> VNode<'a> + 'b,
     {
-        Self {
-            inner: StackNodeStorage::Heap(Box::new(f)),
+        unsafe {
+            let mut ptr: *const _ = &val as &dyn FnOnce(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>();
+
+            if info.len() * mem::size_of::<usize>() + size > mem::size_of::<StackHeapSize>() {
+                log::error!("lazy nodes was too large to fit into stack. falling back to heap");
+
+                Self {
+                    inner: StackNodeStorage::Heap(Box::new(val)),
+                }
+            } else {
+                log::error!("lazy nodes fits on stack!");
+                let mut buf: StackHeapSize = [0; 12];
+
+                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;
+                    }
+                }
+
+                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);
+                }
+
+                Self {
+                    inner: StackNodeStorage::Stack(LazyStack { _align: [], buf }),
+                }
+            }
         }
     }
 
     pub fn call(self, f: NodeFactory<'a>) -> VNode<'a> {
         match self.inner {
             StackNodeStorage::Heap(lazy) => lazy(f),
-            _ => todo!(),
+            StackNodeStorage::Stack(stack) => stack.call(f),
         }
     }
 }
 
-type StackHeapSize = [usize; 12];
+struct LazyStack {
+    _align: [u64; 0],
+    buf: StackHeapSize,
+}
 
-enum StackNodeStorage<'a, 'b> {
-    Stack {
-        next_ofs: usize,
-        buf: StackHeapSize,
-        width: usize,
-    },
-    Heap(Box<dyn FnOnce(NodeFactory<'a>) -> VNode<'a> + 'b>),
+impl LazyStack {
+    unsafe fn create_boxed<'a>(&mut self) -> Box<dyn FnOnce(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 FnOnce(NodeFactory<'a>) -> VNode<'a> =
+            make_fat_ptr(data[..].as_ptr() as usize, &data[info_ofs..]);
+
+        Box::from_raw(g)
+    }
+
+    fn call(mut self, f: NodeFactory) -> VNode {
+        let boxed = unsafe { self.create_boxed() };
+        boxed(f)
+    }
+}
+impl Drop for LazyStack {
+    fn drop(&mut self) {
+        let boxed = unsafe { self.create_boxed() };
+        mem::drop(boxed);
+    }
 }
 
-// impl<'a, F: FnOnce(NodeFactory<'a>) -> VNode<'a>> LazyNodes<'a, F> {
-//     pub fn new(f: F) -> Self {
-//         // let width = std::mem?::size_of::<F>();
-//         // let b: Box<dyn FnOnce(NodeFactory<'a>) -> VNode<'a>> = Box::new(f);
-
-//         todo!()
-//         // Self { inner: b }
-//         // todo!()
-
-//         // if width > std::mem::size_of::<StackHeapSize>() {
-//         //     let g: Box<dyn for<'b> FnOnce(NodeFactory<'b>) -> VNode<'b> + 'g> = Box::new(f);
-//         //     LazyNodes {
-//         //         inner: StackNodeStorage::Heap(g),
-//         //     }
-//         // } else {
-//         //     let mut buf = [0; 12];
-//         //     let mut next_ofs = 0;
-//         //     next_ofs += 1;
-//         //     LazyNodes {
-//         //         inner: StackNodeStorage::Stack {
-//         //             next_ofs,
-//         //             buf,
-//         //             width,
-//         //         },
-//         //     }
-//         // }
-//     }
-// }
-
-// // Our blanket impl
-// impl<'a> IntoIterator for LazyNodes<'a>
-// // where
-// //     F: FnOnce(NodeFactory<'a>) -> VNode<'a>,
-// // impl<'a, F> IntoIterator for LazyNodes<'a, F>
-// // where
-// //     F: FnOnce(NodeFactory<'a>) -> VNode<'a>,
-// {
-//     type Item = Self;
-//     type IntoIter = std::iter::Once<Self::Item>;
-//     fn into_iter(self) -> Self::IntoIter {
-//         std::iter::once(self)
-//     }
-// }
-
-// // Our blanket impl
-// impl IntoVNode for LazyNodes<'_> {
-//     // impl<'a, F: FnOnce(NodeFactory<'a>) -> VNode<'a>> IntoVNode<'a> for LazyNodes<'a, F> {
-//     fn into_vnode<'a>(self, cx: NodeFactory<'a>) -> VNode<'a> {
-//         todo!()
-//         // match self.inner {
-//         //     StackNodeStorage::Stack {
-//         //         buf,
-//         //         next_ofs,
-//         //         width,
-//         //     } => {
-//         //         // get the start of the allocation
-//         //         let r = &buf[0];
-
-//         //         // recast the allocation as dyn FnOnce
-
-//         //         // pretend the FnOnce is box
-//         //         let g: Box<dyn FnOnce(NodeFactory<'a>) -> VNode<'a>> = todo!();
-//         //         // Box::from_raw(r as *const usize as *mut dyn FnOnce(NodeFactory<'a>));
-
-//         //         // use Box's ability to act as FnOnce
-//         //         g(cx)
-//         //     }
-//         //     StackNodeStorage::Heap(b) => b(cx),
-//         // }
-//     }
-// }
+/// Obtain mutable access to a pointer's words
+fn ptr_as_slice<T>(ptr: &mut T) -> &mut [usize] {
+    assert!(mem::size_of::<T>() % mem::size_of::<usize>() == 0);
+    let words = mem::size_of::<T>() / mem::size_of::<usize>();
+    // SAFE: Points to valid memory (a raw pointer)
+    unsafe { core::slice::from_raw_parts_mut(ptr as *mut _ as *mut usize, words) }
+}
+
+/// Re-construct a fat pointer
+unsafe fn make_fat_ptr<T: ?Sized>(data_ptr: usize, meta_vals: &[usize]) -> *mut T {
+    let mut rv = mem::MaybeUninit::<*mut T>::uninit();
+    {
+        let s = ptr_as_slice(&mut rv);
+        s[0] = data_ptr;
+        s[1..].copy_from_slice(meta_vals);
+    }
+    let rv = rv.assume_init();
+    assert_eq!(rv as *const (), data_ptr as *const ());
+    rv
+}
+
+fn round_to_words(len: usize) -> usize {
+    (len + mem::size_of::<usize>() - 1) / mem::size_of::<usize>()
+}