use std::task::{Context, Poll};
use derive_where::derive_where;
use futures_util::FutureExt;
use crate::{
error::{Error, Result},
BoxFuture,
};
const STREAMING_ERR: &str = "cursor state access while streaming";
#[derive_where(Debug; S)]
pub(crate) enum PollState<'a, S, O> {
Idle(S),
Running(#[derive_where(skip)] BoxFuture<'a, (S, O)>),
Polling,
}
impl<'a, S, O> PollState<'a, S, O> {
pub(crate) fn new(state: S) -> Self {
Self::Idle(state)
}
pub(crate) fn state(&self) -> Result<&S> {
match self {
Self::Idle(state) => Ok(state),
_ => Err(Error::internal(STREAMING_ERR)),
}
}
pub(crate) fn state_mut(&mut self) -> Result<&mut S> {
match self {
Self::Idle(state) => Ok(state),
_ => Err(Error::internal(STREAMING_ERR)),
}
}
pub(crate) fn into_state(self) -> Result<S> {
match self {
Self::Idle(state) => Ok(state),
_ => Err(Error::internal(STREAMING_ERR)),
}
}
pub(crate) fn take_state(&mut self) -> Option<S> {
match std::mem::replace(self, Self::Polling) {
Self::Idle(state) => Some(state),
_ => None,
}
}
pub(crate) fn poll_next_step<T>(
&mut self,
cx: &mut Context<'_>,
mut start: impl FnMut(S) -> BoxFuture<'a, (S, O)>,
mut finish: impl FnMut(&mut S, O) -> Option<Result<T>>,
) -> Poll<Option<Result<T>>> {
loop {
match std::mem::replace(self, Self::Polling) {
Self::Idle(state) => {
*self = Self::Running(start(state));
continue;
}
Self::Running(mut future) => match future.poll_unpin(cx) {
Poll::Pending => {
*self = Self::Running(future);
return Poll::Pending;
}
Poll::Ready((mut state, out)) => {
let item = finish(&mut state, out);
*self = Self::Idle(state);
return Poll::Ready(item);
}
},
Self::Polling => {
return Poll::Ready(Some(Err(Error::internal(
"attempt to poll stream already in polling state",
))))
}
}
}
}
}
macro_rules! poll_panic {
($r:expr) => {
$r.unwrap()
};
}
pub(crate) use poll_panic;