generator-light 0.1.0

Light stackless generators based on Futures
Documentation 
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

use core::hint::unreachable_unchecked;
use core::pin::Pin;
use core::task::{Context, Poll};

use crate::GeneratorState;
use crate::{Yielder, fn_trait::FnOnceOutput};

use crate::gen_context::{GeneratorContext, make_yielder_waker};

enum Gen<Init, G> {
    Init(Option<Init>),
    Gen(G),
}

pub(crate) struct Generator<F, Y, R>
where
    F: for<'a> FnOnceOutput<Yielder<'a, Y, R>, R, Out: Future>,
{
    generator: Gen<F, <F as FnOnceOutput<Yielder<'static, Y, R>, R>>::Out>,
}

impl<F, Y, R, Output> Generator<F, Y, R>
where
    F: for<'a> FnOnceOutput<Yielder<'a, Y, R>, R, Out: Future<Output = Output>>,
{
    fn resume_impl(self: Pin<&mut Self>, val: R) -> GeneratorState<Y, Output> {
        // Safety: nor context, nor initialized future will be moved out
        // from this mutable reference
        let this = unsafe { self.get_unchecked_mut() };
        let gen_state: &mut Gen<_, _> = &mut this.generator;
        let state = GeneratorContext::new();
        let waker = unsafe { make_yielder_waker(&state) };
        let mut poll_context = Context::from_waker(&waker);

        let f = match gen_state {
            Gen::Gen(f) => f,
            Gen::Init(init) => {
                // in initial state it's always present
                let init = unsafe { init.take().unwrap_unchecked() };

                let f = init.call(Yielder::new(), val);
                *gen_state = Gen::Gen(f);
                // Poll one with initial seed
                let Gen::Gen(f) = gen_state else {
                    unsafe { unreachable_unchecked() }
                };
                let pinned = unsafe { Pin::new_unchecked(f) };
                return match pinned.poll(&mut poll_context) {
                    Poll::Ready(out) => GeneratorState::Complete(out),
                    Poll::Pending => GeneratorState::Yield(state.take_yielded()),
                };
            }
        };
        // self is pinned.
        // it won't be reinitialized again -> we can safely pin-project here
        let pinned = unsafe { Pin::new_unchecked(f) };
        state.resume(val);
        match pinned.poll(&mut poll_context) {
            Poll::Ready(out) => GeneratorState::Complete(out),
            Poll::Pending => GeneratorState::Yield(state.take_yielded()),
        }
    }

    pub(crate) const fn new(f: F) -> Self {
        Self {
            // ctx: GeneratorContext::new(),
            generator: Gen::Init(Some(f)),
            // _pinned: PhantomPinned,
        }
    }
}

impl<F, Y, R, Output> crate::Generator<R> for Generator<F, Y, R>
where
    F: for<'a> FnOnceOutput<Yielder<'a, Y, R>, R, Out: Future<Output = Output>>,
{
    type Yield = Y;
    type Return = Output;

    fn resume(self: Pin<&mut Self>, resume_val: R) -> GeneratorState<Self::Yield, Self::Return> {
        self.resume_impl(resume_val)
    }
}