📝 only static scope yield is safe (#49)

benches/lib.rs

@@ -131,7 +131,7 @@ fn create_gen(b: &mut Bencher) {
 #[bench]
 fn init_gen(b: &mut Bencher) {
     let clo_gen = || {
-        |mut s: Scope<(), _>| {
+        |mut s: Scope<'_, 'static, (), _>| {
             let mut i = 0;
             loop {
                 match s.yield_(i) {

examples/lifetime.rs

@@ -0,0 +1,23 @@
+fn main() {
+    let str = "foo".to_string();
+
+    let mut gen = generator::Gn::new_scoped(|mut s| {
+        std::thread::scope(|s2| {
+            s2.spawn(|| {
+                std::thread::sleep(std::time::Duration::from_millis(500));
+                println!("{str}");
+            });
+            // here we can't use `yield_` because it still ref to `str`
+            // `yield_` only impl for static captured lifetime
+            // s.yield_(());
+            unsafe { s.yield_unsafe(()) };
+        });
+        generator::done!();
+    });
+
+    gen.next();
+    // std::mem::forget(gen);
+    // drop(gen);
+    // drop(str);
+    std::thread::sleep(std::time::Duration::from_millis(1000));
+}

src/gen_impl.rs

@@ -32,8 +32,9 @@ unsafe impl<A: Send, T: Send> Send for Generator<'static, A, T> {}
 
 impl<'a, A, T> Generator<'a, A, T> {
     /// init a heap based generator with scoped closure
-    pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + Send + 'a>(&mut self, f: F)
+    pub fn scoped_init<F>(&mut self, f: F)
     where
+        for<'scope> F: FnOnce(Scope<'scope, 'a, A, T>) -> T + Send + 'a,
         T: Send + 'a,
         A: Send + 'a,
     {
@@ -55,8 +56,9 @@ pub type LocalGenerator<'a, A, T> = GeneratorObj<'a, A, T, true>;
 
 impl<'a, A, T> LocalGenerator<'a, A, T> {
     /// init a heap based generator with scoped closure
-    pub fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + 'a>(&mut self, f: F)
+    pub fn scoped_init<F>(&mut self, f: F)
     where
+        for<'scope> F: FnOnce(Scope<'scope, 'a, A, T>) -> T + 'a,
         T: 'a,
         A: 'a,
     {
@@ -181,7 +183,7 @@ impl<A> Gn<A> {
     /// create a scoped generator with default stack size
     pub fn new_scoped<'a, T, F>(f: F) -> Generator<'a, A, T>
     where
-        F: FnOnce(Scope<A, T>) -> T + Send + 'a,
+        for<'scope> F: FnOnce(Scope<'scope, 'a, A, T>) -> T + Send + 'a,
         T: Send + 'a,
         A: Send + 'a,
     {
@@ -201,7 +203,7 @@ impl<A> Gn<A> {
     /// create a scoped generator with specified stack size
     pub fn new_scoped_opt<'a, T, F>(size: usize, f: F) -> Generator<'a, A, T>
     where
-        F: FnOnce(Scope<A, T>) -> T + Send + 'a,
+        for<'scope> F: FnOnce(Scope<'scope, 'a, A, T>) -> T + Send + 'a,
         T: Send + 'a,
         A: Send + 'a,
     {
@@ -302,8 +304,9 @@ impl<'a, A, T> GeneratorImpl<'a, A, T> {
     }
 
     /// init a heap based generator with scoped closure
-    fn scoped_init<F: FnOnce(Scope<'a, A, T>) -> T + 'a>(&mut self, f: F)
+    fn scoped_init<F>(&mut self, f: F)
     where
+        for<'scope> F: FnOnce(Scope<'scope, 'a, A, T>) -> T + 'a,
         T: 'a,
         A: 'a,
     {

src/scope.rs

@@ -3,6 +3,7 @@
 //! generator yield implementation
 //!
 
+use std::marker::PhantomData;
 use std::sync::atomic;
 
 use crate::gen_impl::Generator;
@@ -12,15 +13,20 @@ use crate::yield_::raw_yield_now;
 /// passed in scope type
 /// it not use the context to pass data, but keep it's own data ref
 /// this struct provide both compile type info and runtime data
-pub struct Scope<'a, A, T> {
+pub struct Scope<'scope, 'a, A, T> {
     para: &'a mut Option<A>,
     ret: &'a mut Option<T>,
+    scope: PhantomData<&'scope mut &'scope ()>,
 }
 
-impl<'a, A, T> Scope<'a, A, T> {
+impl<'scope, 'a, A, T> Scope<'scope, 'a, A, T> {
     /// create a new scope object
     pub(crate) fn new(para: &'a mut Option<A>, ret: &'a mut Option<T>) -> Self {
-        Scope { para, ret }
+        Scope {
+            para,
+            ret,
+            scope: PhantomData,
+        }
     }
 
     /// set current generator return value
@@ -62,18 +68,22 @@ impl<'a, A, T> Scope<'a, A, T> {
     }
 
     /// yield and get the send para
-    // it's totally safe that we can refer to the function block
-    // since we will come back later
+    /// # Safety
+    /// When yield out, the reference of the captured data must be still valid
+    /// normally, you should always call the `drop` of the generator
     #[inline]
-    pub fn yield_(&mut self, v: T) -> Option<A> {
+    pub unsafe fn yield_unsafe(&mut self, v: T) -> Option<A> {
         self.yield_with(v);
         atomic::compiler_fence(atomic::Ordering::Acquire);
         self.get_yield()
     }
 
     /// `yield_from`
     /// the from generator must has the same type as itself
-    pub fn yield_from(&mut self, mut g: Generator<A, T>) -> Option<A> {
+    /// # Safety
+    /// When yield out, the reference of the captured data must be still valid
+    /// normally, you should always call the `drop` of the generator
+    pub unsafe fn yield_from_unsafe(&mut self, mut g: Generator<A, T>) -> Option<A> {
         let env = ContextStack::current();
         let context = env.top();
         let mut p = self.get_yield();
@@ -88,3 +98,19 @@ impl<'a, A, T> Scope<'a, A, T> {
         p
     }
 }
+
+impl<'scope, A, T> Scope<'scope, 'static, A, T> {
+    /// yield and get the send para
+    // it's totally safe that we can refer to the function block
+    // since we will come back later
+    #[inline]
+    pub fn yield_(&mut self, v: T) -> Option<A> {
+        unsafe { self.yield_unsafe(v) }
+    }
+
+    /// `yield_from`
+    /// the from generator must has the same type as itself
+    pub fn yield_from(&mut self, g: Generator<A, T>) -> Option<A> {
+        unsafe { self.yield_from_unsafe(g) }
+    }
+}

tests/lib.rs

@@ -482,7 +482,7 @@ fn test_scope_yield_from_send() {
 #[test]
 fn test_re_init() {
     let clo = || {
-        |mut s: Scope<(), _>| {
+        |mut s: Scope<'_, 'static, (), _>| {
             s.yield_(0);
             s.yield_(3);
             5