genawaiter/rc/
generator.rs

1use crate::{
2    core::{advance, async_advance, Airlock as _, Next},
3    ops::{Coroutine, GeneratorState},
4    rc::{engine::Airlock, Co},
5};
6use std::{future::Future, pin::Pin};
7
8/// This is a generator which stores its state on the heap.
9///
10/// [_See the module-level docs for examples._](.)
11pub struct Gen<Y, R, F: Future> {
12    airlock: Airlock<Y, R>,
13    future: Pin<Box<F>>,
14}
15
16impl<Y, R, F: Future> Gen<Y, R, F> {
17    /// Creates a new generator from a function.
18    ///
19    /// The function accepts a [`Co`] object, and returns a future. Every time
20    /// the generator is resumed, the future is polled. Each time the future is
21    /// polled, it should do one of two things:
22    ///
23    /// - Call `co.yield_()`, and then return `Poll::Pending`.
24    /// - Drop the `Co`, and then return `Poll::Ready`.
25    ///
26    /// Typically this exchange will happen in the context of an `async fn`.
27    ///
28    /// [_See the module-level docs for examples._](.)
29    pub fn new(producer: impl FnOnce(Co<Y, R>) -> F) -> Self {
30        let airlock = Airlock::default();
31        let future = { Box::pin(producer(Co::new(airlock.clone()))) };
32        Self { airlock, future }
33    }
34
35    /// Resumes execution of the generator.
36    ///
37    /// `arg` is the resume argument. If the generator was previously paused by
38    /// awaiting a future returned from `co.yield()`, that future will complete,
39    /// and return `arg`.
40    ///
41    /// If the generator yields a value, `Yielded` is returned. Otherwise,
42    /// `Completed` is returned.
43    ///
44    /// [_See the module-level docs for examples._](.)
45    pub fn resume_with(&mut self, arg: R) -> GeneratorState<Y, F::Output> {
46        self.airlock.replace(Next::Resume(arg));
47        advance(self.future.as_mut(), &self.airlock)
48    }
49}
50
51impl<Y, F: Future> Gen<Y, (), F> {
52    /// Resumes execution of the generator.
53    ///
54    /// If the generator yields a value, `Yielded` is returned. Otherwise,
55    /// `Completed` is returned.
56    ///
57    /// [_See the module-level docs for examples._](.)
58    pub fn resume(&mut self) -> GeneratorState<Y, F::Output> {
59        self.resume_with(())
60    }
61
62    /// Resumes execution of the generator.
63    ///
64    /// If the generator pauses without yielding, `Poll::Pending` is returned.
65    /// If the generator yields a value, `Poll::Ready(Yielded)` is returned.
66    /// Otherwise, `Poll::Ready(Completed)` is returned.
67    ///
68    /// [_See the module-level docs for examples._](.)
69    pub fn async_resume(
70        &mut self,
71    ) -> impl Future<Output = GeneratorState<Y, F::Output>> + '_ {
72        self.airlock.replace(Next::Resume(()));
73        async_advance(self.future.as_mut(), self.airlock.clone())
74    }
75}
76
77impl<Y, R, F: Future> Coroutine for Gen<Y, R, F> {
78    type Yield = Y;
79    type Resume = R;
80    type Return = F::Output;
81
82    fn resume_with(
83        mut self: Pin<&mut Self>,
84        arg: R,
85    ) -> GeneratorState<Self::Yield, Self::Return> {
86        Self::resume_with(&mut *self, arg)
87    }
88}