Skip to main content

airio_core/upgrade/
select.rs

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