1use std::iter::{Chain, Map};
22
23use either::Either;
24use futures::future;
25
26use crate::{
27 either::EitherFuture,
28 upgrade::{
29 InboundConnectionUpgrade, InboundUpgrade, OutboundConnectionUpgrade, OutboundUpgrade,
30 UpgradeInfo,
31 },
32};
33
34#[derive(Debug, Clone)]
39pub struct SelectUpgrade<A, B>(A, B);
40
41impl<A, B> SelectUpgrade<A, B> {
42 pub fn new(a: A, b: B) -> Self {
46 SelectUpgrade(a, b)
47 }
48}
49
50impl<A, B> UpgradeInfo for SelectUpgrade<A, B>
51where
52 A: UpgradeInfo,
53 B: UpgradeInfo,
54{
55 type Info = Either<A::Info, B::Info>;
56 type InfoIter = Chain<
57 Map<<A::InfoIter as IntoIterator>::IntoIter, fn(A::Info) -> Self::Info>,
58 Map<<B::InfoIter as IntoIterator>::IntoIter, fn(B::Info) -> Self::Info>,
59 >;
60
61 fn protocol_info(&self) -> Self::InfoIter {
62 let a = self
63 .0
64 .protocol_info()
65 .into_iter()
66 .map(Either::Left as fn(A::Info) -> _);
67 let b = self
68 .1
69 .protocol_info()
70 .into_iter()
71 .map(Either::Right as fn(B::Info) -> _);
72
73 a.chain(b)
74 }
75}
76
77impl<C, A, B, TA, TB, EA, EB> InboundUpgrade<C> for SelectUpgrade<A, B>
78where
79 A: InboundUpgrade<C, Output = TA, Error = EA>,
80 B: InboundUpgrade<C, Output = TB, Error = EB>,
81{
82 type Output = future::Either<TA, TB>;
83 type Error = Either<EA, EB>;
84 type Future = EitherFuture<A::Future, B::Future>;
85
86 fn upgrade_inbound(self, sock: C, info: Self::Info) -> Self::Future {
87 match info {
88 Either::Left(info) => EitherFuture::First(self.0.upgrade_inbound(sock, info)),
89 Either::Right(info) => EitherFuture::Second(self.1.upgrade_inbound(sock, info)),
90 }
91 }
92}
93
94impl<C, A, B, TA, TB, EA, EB> InboundConnectionUpgrade<C> for SelectUpgrade<A, B>
95where
96 A: InboundConnectionUpgrade<C, Output = TA, Error = EA>,
97 B: InboundConnectionUpgrade<C, Output = TB, Error = EB>,
98{
99 type Output = future::Either<TA, TB>;
100 type Error = Either<EA, EB>;
101 type Future = EitherFuture<A::Future, B::Future>;
102
103 fn upgrade_inbound(self, sock: C, info: Self::Info) -> Self::Future {
104 match info {
105 Either::Left(info) => EitherFuture::First(self.0.upgrade_inbound(sock, info)),
106 Either::Right(info) => EitherFuture::Second(self.1.upgrade_inbound(sock, info)),
107 }
108 }
109}
110
111impl<C, A, B, TA, TB, EA, EB> OutboundUpgrade<C> for SelectUpgrade<A, B>
112where
113 A: OutboundUpgrade<C, Output = TA, Error = EA>,
114 B: OutboundUpgrade<C, Output = TB, Error = EB>,
115{
116 type Output = future::Either<TA, TB>;
117 type Error = Either<EA, EB>;
118 type Future = EitherFuture<A::Future, B::Future>;
119
120 fn upgrade_outbound(self, sock: C, info: Self::Info) -> Self::Future {
121 match info {
122 Either::Left(info) => EitherFuture::First(self.0.upgrade_outbound(sock, info)),
123 Either::Right(info) => EitherFuture::Second(self.1.upgrade_outbound(sock, info)),
124 }
125 }
126}
127
128impl<C, A, B, TA, TB, EA, EB> OutboundConnectionUpgrade<C> for SelectUpgrade<A, B>
129where
130 A: OutboundConnectionUpgrade<C, Output = TA, Error = EA>,
131 B: OutboundConnectionUpgrade<C, Output = TB, Error = EB>,
132{
133 type Output = future::Either<TA, TB>;
134 type Error = Either<EA, EB>;
135 type Future = EitherFuture<A::Future, B::Future>;
136
137 fn upgrade_outbound(self, sock: C, info: Self::Info) -> Self::Future {
138 match info {
139 Either::Left(info) => EitherFuture::First(self.0.upgrade_outbound(sock, info)),
140 Either::Right(info) => EitherFuture::Second(self.1.upgrade_outbound(sock, info)),
141 }
142 }
143}