use futures_core::{ready, Stream};
use pin_project::{pin_project, project};
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
#[must_use = "futures do nothing unless polled"]
#[pin_project]
pub struct StreamProcessor<S, P, F>
where
F: Future,
{
#[pin]
receiver: S,
processor: P,
#[pin]
state: State<F>,
}
impl<S, P, F> StreamProcessor<S, P, F>
where
F: Future,
{
pub fn new(receiver: S, processor: P) -> Self {
StreamProcessor {
receiver,
processor,
state: State::WaitingForItem,
}
}
}
#[pin_project]
enum State<F> {
WaitingForItem,
RunningProcessor(#[pin] F),
}
impl<S, P, F> Future for StreamProcessor<S, P, F>
where
S: Stream,
P: FnMut(S::Item) -> F,
F: Future,
{
type Output = ();
#[project]
fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
let this = self.as_mut().project();
#[project]
let next = match this.state.project() {
State::WaitingForItem => {
if let Some(item) = ready!(this.receiver.poll_next(ctx)) {
let fut = (this.processor)(item);
State::RunningProcessor(fut)
} else {
return Poll::Ready(());
}
}
State::RunningProcessor(fut) => {
ready!(fut.poll(ctx));
State::WaitingForItem
}
};
self.as_mut().project().state.set(next);
}
}
}