async_gen/
async_iter.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use futures_core::Stream;
7
8use crate::{AsyncGenerator, GeneratorState};
9
10/// An async iterator over the values yielded by an underlying generator.
11///
12/// ## Example
13///
14/// ```
15/// use async_gen::{gen, AsyncIter};
16/// use futures::{StreamExt, Stream};
17///
18/// fn get_async_iter() -> impl Stream<Item = i32> {
19///     AsyncIter::from(gen! {
20///         yield 1;
21///         yield 2;
22///         yield 3;
23///     })
24/// }
25///
26/// #[nio::main]
27/// async fn main() {
28///     let it = get_async_iter();
29///     let v: Vec<_> = it.collect().await;
30///     assert_eq!(v, [1, 2, 3]);
31/// }
32/// ```
33#[derive(Clone)]
34pub struct AsyncIter<G> {
35    gen: G,
36}
37
38impl<G> AsyncIter<G>
39where
40    G: AsyncGenerator<Return = ()>,
41{
42    /// See [`Stream::poll_next`] for more details.
43    #[inline]
44    pub fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<G::Yield>> {
45        unsafe {
46            let me = self.get_unchecked_mut();
47            Pin::new_unchecked(&mut me.gen).poll_resume(cx)
48        }
49        .map(|s| match s {
50            GeneratorState::Yielded(val) => Some(val),
51            GeneratorState::Complete(()) => None,
52        })
53    }
54}
55
56impl<G> Stream for AsyncIter<G>
57where
58    G: AsyncGenerator<Return = ()>,
59{
60    type Item = G::Yield;
61
62    #[inline]
63    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
64        AsyncIter::poll_next(self, cx)
65    }
66}
67
68/// Converts an [`AsyncGenerator`] into an async iterator.
69pub fn async_iter_from<G>(gen: G) -> impl Stream<Item = G::Yield>
70where
71    G: AsyncGenerator<Return = ()>,
72{
73    AsyncIter { gen }
74}
75
76impl<G> From<G> for AsyncIter<G>
77where
78    G: AsyncGenerator<Return = ()>,
79{
80    fn from(gen: G) -> Self {
81        AsyncIter { gen }
82    }
83}
async_gen/async_iter.rs

async_gen/
async_iter.rs
