use std::{marker::PhantomData, pin::Pin};
use futures::{Stream, StreamExt};
pub trait AsyncIterator: Send {
type Item: Send;
fn next(self: Pin<&mut Self>) -> impl futures::Future<Output = Option<Self::Item>> + Send + '_;
fn size_hint(&self) -> (usize, Option<usize>) {
(0, None)
}
fn map<U: Send, C: Send + FnMut(Self::Item) -> U>(self, cb: C) -> impl AsyncIterator<Item = U>
where
Self: Sized,
{
Map { inner: self, cb }
}
}
pin_project_lite::pin_project! {
pub struct Map<U, I: AsyncIterator, F: FnMut(I::Item) ->U> {
#[pin]
inner: I,
cb: F,
}
}
impl<U: Send, I: AsyncIterator, F: Send + FnMut(I::Item) -> U> AsyncIterator for Map<U, I, F> {
type Item = U;
async fn next(self: Pin<&mut Self>) -> Option<U> {
let this = self.project();
Some((this.cb)(this.inner.next().await?))
}
}
pub struct Iter<I: Iterator> {
inner: I,
}
impl<I: Send + Iterator + Unpin> AsyncIterator for Iter<I>
where
I::Item: Send,
{
type Item = I::Item;
#[inline]
async fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
self.get_mut().inner.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
pub fn iter<I: IntoIterator + Send>(inner: I) -> Iter<I::IntoIter>
where
I::IntoIter: Send + Unpin,
I::Item: Send,
{
Iter {
inner: inner.into_iter(),
}
}
pin_project_lite::pin_project! {
pub struct StreamIter<S: Stream> {
#[pin]
inner: S,
}
}
impl<S: Send + Stream> AsyncIterator for StreamIter<S>
where
S::Item: Send,
{
type Item = S::Item;
#[inline]
async fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
self.project().inner.next().await
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
pub fn stream<S: Send + Stream>(inner: S) -> StreamIter<S>
where
S::Item: Send,
{
StreamIter { inner }
}
pub struct Once<I>(Option<I>);
impl<I: Send + Unpin> AsyncIterator for Once<I> {
type Item = I;
async fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
self.get_mut().0.take()
}
}
pub fn once<I: Send + Unpin>(item: I) -> Once<I> {
Once(Some(item))
}
pub struct Empty<I>(PhantomData<I>);
impl<I: Send> AsyncIterator for Empty<I> {
type Item = I;
async fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
None
}
}
pub fn empty<I: Send + Unpin>() -> Empty<I> {
Empty(Default::default())
}
impl<I: Send> AsyncIterator for tokio::sync::mpsc::Receiver<I> {
type Item = I;
async fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
self.get_mut().recv().await
}
}
impl<I: Send> AsyncIterator for tokio::sync::oneshot::Receiver<I> {
type Item = I;
async fn next(self: Pin<&mut Self>) -> Option<Self::Item> {
match self.await {
Ok(item) => Some(item),
Err(_) => None,
}
}
}