1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

//! # yield
//!
//! generator yield implementation
//!

use std::marker::PhantomData;
use std::sync::atomic;

use crate::gen_impl::Generator;
use crate::rt::{Context, ContextStack, Error};
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<'scope, 'a, A, T> {
    para: &'a mut Option<A>,
    ret: &'a mut Option<T>,
    scope: PhantomData<&'scope mut &'scope ()>,
}

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: PhantomData,
        }
    }

    /// set current generator return value
    #[inline]
    fn set_ret(&mut self, v: T) {
        *self.ret = Some(v);
    }

    /// raw yield without catch passed in para
    #[inline]
    fn raw_yield(&mut self, env: &ContextStack, context: &mut Context, v: T) {
        // check the context
        if !context.is_generator() {
            panic!("yield from none generator context");
        }

        self.set_ret(v);
        context._ref -= 1;
        raw_yield_now(env, context);

        // here we just panic to exit the func
        if context._ref != 1 {
            std::panic::panic_any(Error::Cancel);
        }
    }

    /// yield something without catch passed in para
    #[inline]
    pub fn yield_with(&mut self, v: T) {
        let env = ContextStack::current();
        let context = env.top();
        self.raw_yield(&env, context, v);
    }

    /// get current generator send para
    #[inline]
    pub fn get_yield(&mut self) -> Option<A> {
        self.para.take()
    }

    /// yield and get the send para
    /// # 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 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_unsafe`
    /// the from generator must has the same type as itself
    /// # 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();
        while !g.is_done() {
            match g.raw_send(p) {
                None => return None,
                Some(r) => self.raw_yield(&env, context, r),
            }
            p = self.get_yield();
        }
        drop(g); // explicitly consume g
        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) }
    }
}