either_body.rs

use std::backtrace::Backtrace;
use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::pin::Pin;
use std::task::{Context, Poll};

use bytes::Buf;
use http::HeaderMap;
use http_body::{Body, SizeHint};

pub enum Either<L, R> {
    Left(L),
    Right(R),
}

impl<L: Error, R: Error> Debug for Either<L, R> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Either::Left(left) => Debug::fmt(left, f),
            Either::Right(right) => Debug::fmt(right, f),
        }
    }
}

impl<L: Error, R: Error> Display for Either<L, R> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Either::Left(left) => Display::fmt(left, f),
            Either::Right(right) => Display::fmt(right, f),
        }
    }
}

impl<L: Error, R: Error> Error for Either<L, R> {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Either::Left(left) => left.source(),
            Either::Right(right) => right.source(),
        }
    }

    fn backtrace(&self) -> Option<&Backtrace> {
        match self {
            Either::Left(left) => left.backtrace(),
            Either::Right(right) => right.backtrace(),
        }
    }
}

impl<L: Buf, R: Buf> Buf for Either<L, R> {
    fn remaining(&self) -> usize {
        match self {
            Either::Left(left) => left.remaining(),
            Either::Right(right) => right.remaining(),
        }
    }

    fn chunk(&self) -> &[u8] {
        match self {
            Either::Left(left) => left.chunk(),
            Either::Right(right) => right.chunk(),
        }
    }

    fn advance(&mut self, cnt: usize) {
        match self {
            Either::Left(left) => left.advance(cnt),
            Either::Right(right) => right.advance(cnt),
        }
    }
}

#[pin_project::pin_project(project = EitherBodyProj)]
pub enum EitherBody<L, R> {
    Left(#[pin] L),
    Right(#[pin] R),
}

impl<
        L: Body<Data = LBuf, Error = LErr>,
        LBuf: Buf,
        LErr: Error,
        R: Body<Data = RBuf, Error = RErr>,
        RBuf: Buf,
        RErr: Error,
    > Body for EitherBody<L, R>
{
    type Data = Either<LBuf, RBuf>;
    type Error = Either<LErr, RErr>;

    fn poll_data(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
        match self.project() {
            EitherBodyProj::Left(left) => left
                .poll_data(cx)
                .map(|poll| poll.map(|opt| opt.map(Either::Left).map_err(Either::Left))),
            EitherBodyProj::Right(right) => right
                .poll_data(cx)
                .map(|poll| poll.map(|opt| opt.map(Either::Right).map_err(Either::Right))),
        }
    }

    fn poll_trailers(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
        match self.project() {
            EitherBodyProj::Left(left) => left
                .poll_trailers(cx)
                .map(|poll| poll.map_err(Either::Left)),
            EitherBodyProj::Right(right) => right
                .poll_trailers(cx)
                .map(|poll| poll.map_err(Either::Right)),
        }
    }

    fn is_end_stream(&self) -> bool {
        match self {
            Self::Left(left) => left.is_end_stream(),
            Self::Right(right) => right.is_end_stream(),
        }
    }

    fn size_hint(&self) -> SizeHint {
        match self {
            Self::Left(left) => left.size_hint(),
            Self::Right(right) => right.size_hint(),
        }
    }
}