Skip to main content

volans_core/upgrade/
select.rs

1use std::iter::{Chain, Map};
2
3use either::Either;
4use futures::future;
5
6use crate::{
7    either::EitherFuture,
8    upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo},
9};
10
11#[derive(Debug, Clone)]
12pub struct SelectUpgrade<A, B>(A, B);
13
14impl<A, B> SelectUpgrade<A, B> {
15    pub fn new(a: A, b: B) -> Self {
16        SelectUpgrade(a, b)
17    }
18}
19
20impl<A, B> UpgradeInfo for SelectUpgrade<A, B>
21where
22    A: UpgradeInfo,
23    B: UpgradeInfo,
24{
25    type Info = Either<A::Info, B::Info>;
26    type InfoIter = Chain<
27        Map<<A::InfoIter as IntoIterator>::IntoIter, fn(A::Info) -> Self::Info>,
28        Map<<B::InfoIter as IntoIterator>::IntoIter, fn(B::Info) -> Self::Info>,
29    >;
30
31    fn protocol_info(&self) -> Self::InfoIter {
32        let a = self
33            .0
34            .protocol_info()
35            .into_iter()
36            .map(Either::Left as fn(A::Info) -> _);
37        let b = self
38            .1
39            .protocol_info()
40            .into_iter()
41            .map(Either::Right as fn(B::Info) -> _);
42
43        a.chain(b)
44    }
45}
46
47impl<C, A, B, TA, TB, EA, EB> InboundUpgrade<C> for SelectUpgrade<A, B>
48where
49    A: InboundUpgrade<C, Output = TA, Error = EA>,
50    B: InboundUpgrade<C, Output = TB, Error = EB>,
51{
52    type Output = future::Either<TA, TB>;
53    type Error = Either<EA, EB>;
54    type Future = EitherFuture<A::Future, B::Future>;
55
56    fn upgrade_inbound(self, stream: C, info: Self::Info) -> Self::Future {
57        match info {
58            Either::Left(info) => EitherFuture::First(self.0.upgrade_inbound(stream, info)),
59            Either::Right(info) => EitherFuture::Second(self.1.upgrade_inbound(stream, info)),
60        }
61    }
62}
63
64impl<C, A, B, TA, TB, EA, EB> OutboundUpgrade<C> for SelectUpgrade<A, B>
65where
66    A: OutboundUpgrade<C, Output = TA, Error = EA>,
67    B: OutboundUpgrade<C, Output = TB, Error = EB>,
68{
69    type Output = future::Either<TA, TB>;
70    type Error = Either<EA, EB>;
71    type Future = EitherFuture<A::Future, B::Future>;
72
73    fn upgrade_outbound(self, stream: C, info: Self::Info) -> Self::Future {
74        match info {
75            Either::Left(info) => EitherFuture::First(self.0.upgrade_outbound(stream, info)),
76            Either::Right(info) => EitherFuture::Second(self.1.upgrade_outbound(stream, info)),
77        }
78    }
79}