📝 refine StackBox impl

Author: Xudong-Huang

src/gen_impl.rs

@@ -6,7 +6,6 @@
 use std::any::Any;
 use std::fmt;
 use std::marker::PhantomData;
-use std::mem::ManuallyDrop;
 use std::panic;
 use std::thread;
 
@@ -22,8 +21,7 @@ pub const DEFAULT_STACK_SIZE: usize = 0x1000;
 
 /// the generator type
 pub struct Generator<'a, A, T> {
-    _stack: StackBox<Stack>,
-    gen: ManuallyDrop<StackBox<GeneratorImpl<'a, A, T>>>,
+    gen: StackBox<GeneratorImpl<'a, A, T>>,
 }
 
 unsafe impl<A, T> Send for Generator<'static, A, T> {}
@@ -38,11 +36,8 @@ impl<'a, A, T> Generator<'a, A, T> {
     /// function is called twice on the same raw pointer.
     #[inline]
     pub unsafe fn from_raw(raw: *mut usize) -> Self {
-        let g = StackBox::from_raw(raw as *mut GeneratorImpl<'a, A, T>);
-        let stack_ptr = raw.offset(g.size() as isize + 2);
         Generator {
-            _stack: StackBox::from_raw(stack_ptr as *mut Stack),
-            gen: ManuallyDrop::new(g),
+            gen: StackBox::from_raw(raw as *mut GeneratorImpl<'a, A, T>),
         }
     }
 
@@ -95,12 +90,6 @@ impl<'a, A, T> fmt::Debug for Generator<'a, A, T> {
     }
 }
 
-impl<'a, A, T> Drop for Generator<'a, A, T> {
-    fn drop(&mut self) {
-        unsafe { ManuallyDrop::drop(&mut self.gen) }
-    }
-}
-
 /// Generator helper
 pub struct Gn<A = ()> {
     dummy: PhantomData<A>,
@@ -124,13 +113,9 @@ impl<A> Gn<A> {
         T: 'a,
         A: 'a,
     {
-        let mut stack = Stack::new(size);
-        let mut g = GeneratorImpl::<A, T>::new(&mut stack);
-        g.scoped_init(f);
-        Generator {
-            _stack: stack,
-            gen: ManuallyDrop::new(g),
-        }
+        let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
+        gen.scoped_init(f);
+        Generator { gen }
     }
 }
 
@@ -151,14 +136,10 @@ impl<A: Any> Gn<A> {
     where
         F: FnOnce() -> T + 'a,
     {
-        let mut stack = Stack::new(size);
-        let mut g = GeneratorImpl::<A, T>::new(&mut stack);
-        g.init_context();
-        g.init_code(f);
-        Generator {
-            _stack: stack,
-            gen: ManuallyDrop::new(g),
-        }
+        let mut gen = GeneratorImpl::<A, T>::new(Stack::new(size));
+        gen.init_context();
+        gen.init_code(f);
+        Generator { gen }
     }
 }
 
@@ -167,6 +148,8 @@ impl<A: Any> Gn<A> {
 pub struct GeneratorImpl<'a, A, T> {
     // run time context
     context: Context,
+    // stack
+    stack: Stack,
     // save the input
     para: Option<A>,
     // save the output
@@ -192,17 +175,19 @@ impl<'a, A: Any, T: Any> GeneratorImpl<'a, A, T> {
 
 impl<'a, A, T> GeneratorImpl<'a, A, T> {
     /// create a new generator with specified stack size
-    fn new(stack: &mut Stack) -> StackBox<Self> {
+    fn new(mut stack: Stack) -> StackBox<Self> {
+        // the stack box would finally dealloc the stack!
         unsafe {
             let mut stack_box = stack
                 .alloc_uninit_box::<GeneratorImpl<'a, A, T>>()
                 .assume_init();
 
             stack_box.init(GeneratorImpl {
                 para: None,
+                stack,
                 ret: None,
                 f: None,
-                context: Context::new(stack),
+                context: Context::new(),
                 phantom: PhantomData,
             });
             stack_box
@@ -244,11 +229,10 @@ impl<'a, A, T> GeneratorImpl<'a, A, T> {
 
         // init the ref to 0 means that it's ready to start
         self.context._ref = 0;
-        let stack = unsafe { &mut *self.context.stack };
         let ret = &mut self.ret as *mut _;
         let context = &mut self.context as *mut Context;
         // alloc the function on stack
-        let func = StackBox::new_fn_once(stack, move || {
+        let func = StackBox::new_fn_once(&mut self.stack, move || {
             let r = f();
             let ret = unsafe { &mut *ret };
             let _ref = unsafe { (*context)._ref };
@@ -262,10 +246,12 @@ impl<'a, A, T> GeneratorImpl<'a, A, T> {
 
         self.f = Some(func);
 
-        let stk = unsafe { &*self.context.stack };
-        self.context
-            .regs
-            .init_with(gen_init, 0, &mut self.f as *mut _ as *mut usize, stk);
+        self.context.regs.init_with(
+            gen_init,
+            0,
+            &mut self.f as *mut _ as *mut usize,
+            &self.stack,
+        );
     }
 
     /// resume the generator
@@ -411,8 +397,7 @@ impl<'a, A, T> GeneratorImpl<'a, A, T> {
 
     /// get stack total size and used size in word
     pub fn stack_usage(&self) -> (usize, usize) {
-        let stack = unsafe { &*self.context.stack };
-        (stack.size(), stack.get_used_size())
+        (self.stack.size(), self.stack.get_used_size())
     }
 }
 

src/rt.rs

@@ -7,12 +7,11 @@ use std::mem::MaybeUninit;
 use std::ptr;
 
 use crate::reg_context::RegContext;
-use crate::stack::Stack;
 
 thread_local!(
     /// each thread has it's own generator context stack
     static ROOT_CONTEXT: Box<Context> = {
-        let mut root = Box::new(Context::empty());
+        let mut root = Box::new(Context::new());
         let p = &mut *root as *mut _;
         root.parent = p; // init top to current
         root
@@ -50,8 +49,6 @@ pub struct Context {
     child: *mut Context,
     /// parent context
     pub parent: *mut Context,
-    /// generator execution stack
-    pub stack: *mut Stack,
     /// passed in para for send
     pub para: MaybeUninit<*mut dyn Any>,
     /// this is just a buffer for the return value
@@ -65,26 +62,10 @@ pub struct Context {
 }
 
 impl Context {
-    // create for root empty context
-    fn empty() -> Self {
-        Context {
-            regs: RegContext::empty(),
-            stack: ptr::null_mut(),
-            para: MaybeUninit::zeroed(),
-            ret: MaybeUninit::zeroed(),
-            _ref: 1, // none zero means it's not running
-            err: None,
-            child: ptr::null_mut(),
-            parent: ptr::null_mut(),
-            local_data: ptr::null_mut(),
-        }
-    }
-
     /// return a default generator context
-    pub fn new(stack: &mut Stack) -> Context {
+    pub fn new() -> Context {
         Context {
             regs: RegContext::empty(),
-            stack,
             para: MaybeUninit::zeroed(),
             ret: MaybeUninit::zeroed(),
             _ref: 1, // none zero means it's not running

src/stack/mod.rs

@@ -22,10 +22,12 @@ const ALIGN: usize = std::mem::size_of::<StackBoxHeader>();
 const HEADER_SIZE: usize = std::mem::size_of::<StackBoxHeader>() / std::mem::size_of::<usize>();
 
 struct StackBoxHeader {
-    // track the stack offset, saved on stack
-    offset: *mut usize,
+    // track the stack
+    stack: Stack,
     // track how big the data is (in usize)
     data_size: usize,
+    // non zero dealloc the stack
+    need_drop: usize,
 }
 
 /// A pointer type for stack allocation.
@@ -36,7 +38,7 @@ pub struct StackBox<T> {
 
 impl<T> StackBox<T> {
     /// create uninit stack box
-    fn new_unint(stack: &mut Stack) -> MaybeUninit<Self> {
+    fn new_unint(stack: &mut Stack, need_drop: usize) -> MaybeUninit<Self> {
         let offset = unsafe { &mut *stack.get_offset() };
         // alloc the data
         let layout = std::alloc::Layout::new::<T>();
@@ -54,7 +56,8 @@ impl<T> StackBox<T> {
             let mut header = ptr::NonNull::new_unchecked(stack.end() as *mut StackBoxHeader);
             let header = header.as_mut();
             header.data_size = data_size;
-            header.offset = offset;
+            header.need_drop = need_drop;
+            header.stack = stack.shadow_clone();
             std::mem::MaybeUninit::new(StackBox { ptr })
         }
     }
@@ -76,11 +79,6 @@ impl<T> StackBox<T> {
         self.ptr.as_ptr()
     }
 
-    // get the stack ptr
-    pub(crate) fn size(&self) -> usize {
-        self.get_header().data_size
-    }
-
     /// Constructs a StackBox from a raw pointer.
     ///
     /// # Safety
@@ -125,7 +123,7 @@ impl Drop for Func {
         if !self.data.is_null() {
             (self.drop)(self.data);
         }
-        unsafe { *self.offset -= self.size }
+        unsafe { *self.offset -= self.size };
     }
 }
 
@@ -148,13 +146,13 @@ impl<F: FnOnce()> StackBox<F> {
     /// create a functor on the stack
     pub(crate) fn new_fn_once(stack: &mut Stack, data: F) -> Func {
         unsafe {
-            let mut d = Self::new_unint(stack).assume_init();
+            let mut d = Self::new_unint(stack, 0).assume_init();
             d.init(data);
             let header = d.get_header();
             let f = Func {
                 data: d.ptr.as_ptr() as *mut (),
                 size: header.data_size + HEADER_SIZE,
-                offset: header.offset,
+                offset: stack.get_offset(),
                 func: Self::call_once,
                 drop: Self::drop_inner,
             };
@@ -182,8 +180,11 @@ impl<T> Drop for StackBox<T> {
     fn drop(&mut self) {
         let header = self.get_header();
         unsafe {
-            *header.offset -= header.data_size + HEADER_SIZE;
+            *header.stack.get_offset() -= header.data_size + HEADER_SIZE;
             ptr::drop_in_place(self.ptr.as_ptr());
+            if header.need_drop != 0 {
+                header.stack.drop_stack();
+            }
         }
     }
 }
@@ -301,13 +302,15 @@ impl SysStack {
 unsafe impl Send for SysStack {}
 
 /// generator stack
+/// this struct will not dealloc the memroy
+/// instead StackBox<> would track it's usage and dealloc it
 pub struct Stack {
     buf: SysStack,
 }
 
 impl Stack {
     /// Allocate a new stack of `size`. If size = 0, this is a `dummy_stack`
-    pub fn new(size: usize) -> StackBox<Stack> {
+    pub fn new(size: usize) -> Stack {
         let track = (size & 1) != 0;
         let mut bytes = size * std::mem::size_of::<usize>();
         // the minimal size
@@ -319,7 +322,7 @@ impl Stack {
 
         let buf = SysStack::allocate(bytes, true).expect("failed to alloc sys stack");
 
-        let mut stk = Stack { buf };
+        let stk = Stack { buf };
 
         // if size is not even we do the full foot print test
         let count = if track {
@@ -333,15 +336,12 @@ impl Stack {
             let buf = stk.buf.bottom as *mut usize;
             ptr::write_bytes(buf, 0xEE, count);
         }
+
         // init the stack box usage
         let offset = stk.get_offset();
         unsafe { *offset = 1 };
 
-        unsafe {
-            let mut stack = stk.alloc_uninit_box::<Stack>().assume_init();
-            stack.init(stk);
-            stack
-        }
+        stk
     }
 
     /// get used stack size
@@ -380,24 +380,17 @@ impl Stack {
 
     /// alloc buffer on this stack
     pub fn alloc_uninit_box<T>(&mut self) -> MaybeUninit<StackBox<T>> {
-        StackBox::<T>::new_unint(self)
+        // the first obj should set need drop to non zero
+        StackBox::<T>::new_unint(self, 1)
     }
 
     // get offset
     fn get_offset(&self) -> *mut usize {
         unsafe { (self.buf.top as *mut usize).offset(-1) }
     }
-}
 
-impl fmt::Debug for Stack {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let offset = self.get_offset();
-        write!(f, "Statck<{:?}, Offset={}>", self.buf, unsafe { *offset })
-    }
-}
-
-impl Drop for Stack {
-    fn drop(&mut self) {
+    // dealloc the statck
+    fn drop_stack(&self) {
         if self.buf.len() == 0 {
             return;
         }
@@ -408,4 +401,20 @@ impl Drop for Stack {
             sys::deallocate_stack(guard, size_with_guard);
         }
     }
+
+    fn shadow_clone(&self) -> Self {
+        Stack {
+            buf: SysStack {
+                top: self.buf.top,
+                bottom: self.buf.bottom,
+            },
+        }
+    }
+}
+
+impl fmt::Debug for Stack {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let offset = self.get_offset();
+        write!(f, "Statck<{:?}, Offset={}>", self.buf, unsafe { *offset })
+    }
 }