use alloc::{rc::Rc, sync::Arc};
use core::{
future::Future,
marker::PhantomData,
pin::Pin,
task::{Context, Poll},
};
use futures_core::ready;
use pin_project_lite::pin_project;
use crate::{IntoServiceFactory, Service, ServiceFactory};
pub fn apply<T, S, I, Req>(t: T, factory: I) -> ApplyTransform<T, S, Req>
where
I: IntoServiceFactory<S, Req>,
S: ServiceFactory<Req>,
T: Transform<S::Service, Req, InitError = S::InitError>,
{
ApplyTransform::new(t, factory.into_factory())
}
pub trait Transform<S, Req> {
type Response;
type Error;
type Transform: Service<Req, Response = Self::Response, Error = Self::Error>;
type InitError;
type Future: Future<Output = Result<Self::Transform, Self::InitError>>;
fn new_transform(&self, service: S) -> Self::Future;
}
impl<T, S, Req> Transform<S, Req> for Rc<T>
where
T: Transform<S, Req>,
{
type Response = T::Response;
type Error = T::Error;
type Transform = T::Transform;
type InitError = T::InitError;
type Future = T::Future;
fn new_transform(&self, service: S) -> T::Future {
self.as_ref().new_transform(service)
}
}
impl<T, S, Req> Transform<S, Req> for Arc<T>
where
T: Transform<S, Req>,
{
type Response = T::Response;
type Error = T::Error;
type Transform = T::Transform;
type InitError = T::InitError;
type Future = T::Future;
fn new_transform(&self, service: S) -> T::Future {
self.as_ref().new_transform(service)
}
}
pub struct ApplyTransform<T, S, Req>(Rc<(T, S)>, PhantomData<Req>);
impl<T, S, Req> ApplyTransform<T, S, Req>
where
S: ServiceFactory<Req>,
T: Transform<S::Service, Req, InitError = S::InitError>,
{
fn new(t: T, service: S) -> Self {
Self(Rc::new((t, service)), PhantomData)
}
}
impl<T, S, Req> Clone for ApplyTransform<T, S, Req> {
fn clone(&self) -> Self {
ApplyTransform(self.0.clone(), PhantomData)
}
}
impl<T, S, Req> ServiceFactory<Req> for ApplyTransform<T, S, Req>
where
S: ServiceFactory<Req>,
T: Transform<S::Service, Req, InitError = S::InitError>,
{
type Response = T::Response;
type Error = T::Error;
type Config = S::Config;
type Service = T::Transform;
type InitError = T::InitError;
type Future = ApplyTransformFuture<T, S, Req>;
fn new_service(&self, cfg: S::Config) -> Self::Future {
ApplyTransformFuture {
store: self.0.clone(),
state: ApplyTransformFutureState::A {
fut: self.0.as_ref().1.new_service(cfg),
},
}
}
}
pin_project! {
pub struct ApplyTransformFuture<T, S, Req>
where
S: ServiceFactory<Req>,
T: Transform<S::Service, Req, InitError = S::InitError>,
{
store: Rc<(T, S)>,
#[pin]
state: ApplyTransformFutureState<T, S, Req>,
}
}
pin_project! {
#[project = ApplyTransformFutureStateProj]
pub enum ApplyTransformFutureState<T, S, Req>
where
S: ServiceFactory<Req>,
T: Transform<S::Service, Req, InitError = S::InitError>,
{
A { #[pin] fut: S::Future },
B { #[pin] fut: T::Future },
}
}
impl<T, S, Req> Future for ApplyTransformFuture<T, S, Req>
where
S: ServiceFactory<Req>,
T: Transform<S::Service, Req, InitError = S::InitError>,
{
type Output = Result<T::Transform, T::InitError>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.as_mut().project();
match this.state.as_mut().project() {
ApplyTransformFutureStateProj::A { fut } => {
let srv = ready!(fut.poll(cx))?;
let fut = this.store.0.new_transform(srv);
this.state.set(ApplyTransformFutureState::B { fut });
self.poll(cx)
}
ApplyTransformFutureStateProj::B { fut } => fut.poll(cx),
}
}
}