📝 impl generator into_raw/from_raw

Author: Xudong-Huang

src/gen_impl.rs

@@ -26,6 +26,35 @@ pub struct Generator<'a, A, T> {
     gen: ManuallyDrop<StackBox<GeneratorImpl<'a, A, T>>>,
 }
 
+unsafe impl<A, T> Send for Generator<'static, A, T> {}
+
+impl<'a, A, T> Generator<'a, A, T> {
+    /// Constructs a Generator from a raw pointer.
+    ///
+    /// # Safety
+    ///
+    /// This function is unsafe because improper use may lead to
+    /// memory problems. For example, a double-free may occur if the
+    /// 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),
+        }
+    }
+
+    /// Consumes the `Generator`, returning a wrapped raw pointer.
+    #[inline]
+    pub fn into_raw(g: Generator<'a, A, T>) -> *mut usize {
+        let ret = g.gen.as_ptr() as *mut usize;
+        std::mem::forget(g);
+        ret
+    }
+}
+
 impl<'a, A, T> std::ops::Deref for Generator<'a, A, T> {
     type Target = GeneratorImpl<'a, A, T>;
 
@@ -66,7 +95,7 @@ impl<'a, A, T> fmt::Debug for Generator<'a, A, T> {
     }
 }
 
-impl<'a, A, T> std::ops::Drop 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) }
     }
@@ -99,8 +128,8 @@ impl<A> Gn<A> {
         let mut g = GeneratorImpl::<A, T>::new(&mut stack);
         g.scoped_init(f);
         Generator {
-            gen: ManuallyDrop::new(g),
             _stack: stack,
+            gen: ManuallyDrop::new(g),
         }
     }
 }
@@ -127,8 +156,8 @@ impl<A: Any> Gn<A> {
         g.init_context();
         g.init_code(f);
         Generator {
-            gen: ManuallyDrop::new(g),
             _stack: stack,
+            gen: ManuallyDrop::new(g),
         }
     }
 }
@@ -148,8 +177,6 @@ pub struct GeneratorImpl<'a, A, T> {
     phantom: PhantomData<&'a T>,
 }
 
-unsafe impl<A, T> Send for GeneratorImpl<'static, A, T> {}
-
 impl<'a, A: Any, T: Any> GeneratorImpl<'a, A, T> {
     /// create a new generator with default stack size
     fn init_context(&mut self) {

src/lib.rs

@@ -26,7 +26,6 @@ mod yield_;
 pub use crate::gen_impl::{Generator, Gn};
 pub use crate::rt::{get_local_data, is_generator, Error};
 pub use crate::scope::Scope;
-pub use crate::stack::StackBox;
 pub use crate::yield_::{
     co_get_yield, co_set_para, co_yield_with, done, get_yield, yield_, yield_from, yield_with,
 };

src/stack/mod.rs

@@ -71,6 +71,16 @@ impl<T> StackBox<T> {
         ptr::write(self.ptr.as_ptr(), data);
     }
 
+    // get the stack ptr
+    pub(crate) fn as_ptr(&self) -> *mut 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
@@ -79,40 +89,40 @@ impl<T> StackBox<T> {
     /// memory problems. For example, a double-free may occur if the
     /// function is called twice on the same raw pointer.
     #[inline]
-    pub unsafe fn from_raw(raw: *mut T) -> Self {
+    pub(crate) unsafe fn from_raw(raw: *mut T) -> Self {
         StackBox {
             ptr: ptr::NonNull::new_unchecked(raw),
         }
     }
 
-    /// Consumes the `StackBox`, returning a wrapped raw pointer.
-    #[inline]
-    pub fn into_raw(b: StackBox<T>) -> *mut T {
-        let ret = b.ptr.as_ptr();
-        std::mem::forget(b);
-        ret
-    }
+    // Consumes the `StackBox`, returning a wrapped raw pointer.
+    // #[inline]
+    // pub(crate) fn into_raw(b: StackBox<T>) -> *mut T {
+    //     let ret = b.ptr.as_ptr();
+    //     std::mem::forget(b);
+    //     ret
+    // }
 }
 
 pub struct Func {
     data: *mut (),
     size: usize,
     offset: *mut usize,
-    is_called: bool,
     func: fn(*mut ()),
     drop: fn(*mut ()),
 }
 
 impl Func {
     pub fn call_once(mut self) {
-        (self.func)(self.data);
-        self.is_called = true;
+        let data = self.data;
+        self.data = ptr::null_mut();
+        (self.func)(data);
     }
 }
 
 impl Drop for Func {
     fn drop(&mut self) {
-        if !self.is_called {
+        if !self.data.is_null() {
             (self.drop)(self.data);
         }
         unsafe { *self.offset -= self.size }
@@ -145,7 +155,6 @@ impl<F: FnOnce()> StackBox<F> {
                 data: d.ptr.as_ptr() as *mut (),
                 size: header.data_size + HEADER_SIZE,
                 offset: header.offset,
-                is_called: false,
                 func: Self::call_once,
                 drop: Self::drop_inner,
             };