use core::{
ops::ControlFlow,
pin::Pin,
task::{Context, Poll},
};
use futures::Future;
use pin_project::pin_project;
use crate::effect::{Effective, ErasedEffective, Ss};
use super::{Async, BoxOrReady, Then};
#[pin_project(project = MapStateProj, project_replace = MapStateReplace)]
enum MapState<Fut, F> {
Incomplete {
#[pin]
future: Fut,
f: F,
},
Completed,
}
#[pin_project]
pub struct Map<Fut, F> {
#[pin]
state: MapState<Fut, F>,
}
impl<Fut, F> Map<Fut, F> {
pub(super) fn new(future: Fut, f: F) -> Self {
Self {
state: MapState::Incomplete { future, f },
}
}
}
impl<'a, Fut, F, T> Future for Map<Fut, F>
where
Fut: Future,
F: FnOnce(Fut::Output) -> T,
{
type Output = T;
#[inline(always)]
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.project();
match this.state.as_mut().project() {
MapStateProj::Incomplete { future, .. } => {
let output = match future.poll(cx) {
Poll::Ready(value) => value,
Poll::Pending => return Poll::Pending,
};
match this.state.project_replace(MapState::Completed) {
MapStateReplace::Incomplete { f, .. } => Poll::Ready(f(output)),
_ => unreachable!(),
}
}
MapStateProj::Completed => {
panic!("Map must not be polled after it returned `Poll::Ready`")
}
}
}
}
impl<'c, 'lt: 'c, Fut0: Ss + 'lt, F0: Ss + 'c, R0: Ss + 'c> Effective<'c> for Map<Fut0, F0>
where
F0: FnOnce(Fut0::Output) -> R0,
Fut0: Future,
Fut0::Output: 'lt,
{
fn into_erased<B>(self) -> ErasedEffective<'c, Self::Output, Self::Effect, B> {
BoxOrReady::Boxed(Box::pin(self))
}
type Effect = Async;
type Output = R0;
type IntoFuture = Self;
fn into_future(self) -> Self::IntoFuture {
self
}
type Loop<'ctx: 'a, 'a, T: Ss + 'a, F: Ss + 'a>
= BoxOrReady<'a, (R0, T)>
where
F: for<'b> FnMut(&'b mut Self::Output) -> ErasedEffective<'b, ControlFlow<T>, Self::Effect, (&'b mut Self::Output, &'ctx ())>,
'c: 'a;
fn r#loop<'ctx: 'a, 'a, T: Ss + 'a, F: Ss + 'a>(self, mut cb: F) -> Self::Loop<'ctx, 'a, T, F>
where
F: for<'b> FnMut(
&'b mut Self::Output,
) -> ErasedEffective<
'b,
ControlFlow<T>,
Self::Effect,
(&'b mut Self::Output, &'ctx ()),
>,
'c: 'a,
{
BoxOrReady::Boxed(Box::pin(async move {
let mut this = self.into_future().await;
loop {
if let ControlFlow::Break(value) = cb(&mut this).into_future().await {
return (this, value);
}
}
}))
}
type Map<'a, T: Ss + 'a, F: Ss + 'a> =
Map<Self, F>
where
F: FnOnce(Self::Output) -> T,
'c: 'a;
fn map<'a, T: Ss + 'a, F: Ss + 'a>(self, cb: F) -> Self::Map<'a, T, F>
where
F: FnOnce(Self::Output) -> T,
'c: 'a,
{
Map::new(self, cb)
}
type Then<'a, T: Ss + 'a, V: Ss + 'a, F: Ss + 'a>
= Then<'a, Self, F, V>
where
F: FnOnce(Self::Output) -> V,
V: Effective<'a, Output = T, Effect = Self::Effect>,
'c: 'a;
fn then<'a, T: Ss + 'a, V: Ss + 'a, F: Ss + 'a>(self, cb: F) -> Self::Then<'a, T, V, F>
where
F: FnOnce(Self::Output) -> V,
V: Effective<'a, Output = T, Effect = Self::Effect>,
'c: 'a,
{
Then::new(self, cb)
}
type AsCtx<'ctx: 'a, 'a, T: Ss + 'a, F: Ss + 'a>
= BoxOrReady<'a, (Self::Output, T)>
where
F: for<'b> FnOnce(&'b mut Self::Output) -> ErasedEffective<'b, T, Self::Effect>,
'c: 'a;
fn as_ctx<'ctx: 'a, 'a, T: Ss + 'a, F: Ss + 'a>(self, cb: F) -> Self::AsCtx<'ctx, 'a, T, F>
where
F: for<'b> FnOnce(&'b mut Self::Output) -> ErasedEffective<'b, T, Self::Effect>,
'c: 'a,
{
BoxOrReady::Boxed(Box::pin(async {
let mut this = self.await;
let result = cb(&mut this).into_future().await;
(this, result)
}))
}
}