use crate::{Error, ErrorCompat, IntoError};
use core::{
future::Future,
marker::PhantomData,
pin::Pin,
task::{Context as TaskContext, Poll},
};
use futures_core_crate::future::TryFuture;
use pin_project::pin_project;
pub trait TryFutureExt: TryFuture + Sized {
fn context<C, E>(self, context: C) -> Context<Self, C, E>
where
C: IntoError<E, Source = Self::Error>,
E: Error + ErrorCompat;
fn with_context<F, C, E>(self, context: F) -> WithContext<Self, F, E>
where
F: FnOnce() -> C,
C: IntoError<E, Source = Self::Error>,
E: Error + ErrorCompat;
}
impl<Fut> TryFutureExt for Fut
where
Fut: TryFuture,
{
fn context<C, E>(self, context: C) -> Context<Self, C, E>
where
C: IntoError<E, Source = Self::Error>,
E: Error + ErrorCompat,
{
Context {
inner: self,
context: Some(context),
_e: PhantomData,
}
}
fn with_context<F, C, E>(self, context: F) -> WithContext<Self, F, E>
where
F: FnOnce() -> C,
C: IntoError<E, Source = Self::Error>,
E: Error + ErrorCompat,
{
WithContext {
inner: self,
context: Some(context),
_e: PhantomData,
}
}
}
#[pin_project]
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Context<Fut, C, E> {
#[pin]
inner: Fut,
context: Option<C>,
_e: PhantomData<E>,
}
impl<Fut, C, E> Future for Context<Fut, C, E>
where
Fut: TryFuture,
C: IntoError<E, Source = Fut::Error>,
E: Error + ErrorCompat,
{
type Output = Result<Fut::Ok, E>;
fn poll(self: Pin<&mut Self>, ctx: &mut TaskContext) -> Poll<Self::Output> {
let this = self.project();
let inner = this.inner;
let context = this.context;
inner.try_poll(ctx).map_err(|error| {
context
.take()
.expect("Cannot poll Context after it resolves")
.into_error(error)
})
}
}
#[pin_project]
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct WithContext<Fut, F, E> {
#[pin]
inner: Fut,
context: Option<F>,
_e: PhantomData<E>,
}
impl<Fut, F, C, E> Future for WithContext<Fut, F, E>
where
Fut: TryFuture,
F: FnOnce() -> C,
C: IntoError<E, Source = Fut::Error>,
E: Error + ErrorCompat,
{
type Output = Result<Fut::Ok, E>;
fn poll(self: Pin<&mut Self>, ctx: &mut TaskContext) -> Poll<Self::Output> {
let this = self.project();
let inner = this.inner;
let context = this.context;
inner.try_poll(ctx).map_err(|error| {
let context = context
.take()
.expect("Cannot poll WithContext after it resolves");
context().into_error(error)
})
}
}