libp2p_core/upgrade/
select.rs

1// Copyright 2018 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21use crate::{
22    either::{EitherOutput, EitherError, EitherFuture2, EitherName},
23    upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo}
24};
25
26/// Upgrade that combines two upgrades into one. Supports all the protocols supported by either
27/// sub-upgrade.
28///
29/// The protocols supported by the first element have a higher priority.
30#[derive(Debug, Clone)]
31pub struct SelectUpgrade<A, B>(A, B);
32
33impl<A, B> SelectUpgrade<A, B> {
34    /// Combines two upgrades into an `SelectUpgrade`.
35    ///
36    /// The protocols supported by the first element have a higher priority.
37    pub fn new(a: A, b: B) -> Self {
38        SelectUpgrade(a, b)
39    }
40}
41
42impl<A, B> UpgradeInfo for SelectUpgrade<A, B>
43where
44    A: UpgradeInfo,
45    B: UpgradeInfo
46{
47    type Info = EitherName<A::Info, B::Info>;
48    type InfoIter = InfoIterChain<
49        <A::InfoIter as IntoIterator>::IntoIter,
50        <B::InfoIter as IntoIterator>::IntoIter
51    >;
52
53    fn protocol_info(&self) -> Self::InfoIter {
54        InfoIterChain(self.0.protocol_info().into_iter(), self.1.protocol_info().into_iter())
55    }
56}
57
58impl<C, A, B, TA, TB, EA, EB> InboundUpgrade<C> for SelectUpgrade<A, B>
59where
60    A: InboundUpgrade<C, Output = TA, Error = EA>,
61    B: InboundUpgrade<C, Output = TB, Error = EB>,
62{
63    type Output = EitherOutput<TA, TB>;
64    type Error = EitherError<EA, EB>;
65    type Future = EitherFuture2<A::Future, B::Future>;
66
67    fn upgrade_inbound(self, sock: C, info: Self::Info) -> Self::Future {
68        match info {
69            EitherName::A(info) => EitherFuture2::A(self.0.upgrade_inbound(sock, info)),
70            EitherName::B(info) => EitherFuture2::B(self.1.upgrade_inbound(sock, info))
71        }
72    }
73}
74
75impl<C, A, B, TA, TB, EA, EB> OutboundUpgrade<C> for SelectUpgrade<A, B>
76where
77    A: OutboundUpgrade<C, Output = TA, Error = EA>,
78    B: OutboundUpgrade<C, Output = TB, Error = EB>,
79{
80    type Output = EitherOutput<TA, TB>;
81    type Error = EitherError<EA, EB>;
82    type Future = EitherFuture2<A::Future, B::Future>;
83
84    fn upgrade_outbound(self, sock: C, info: Self::Info) -> Self::Future {
85        match info {
86            EitherName::A(info) => EitherFuture2::A(self.0.upgrade_outbound(sock, info)),
87            EitherName::B(info) => EitherFuture2::B(self.1.upgrade_outbound(sock, info))
88        }
89    }
90}
91
92/// Iterator that combines the protocol names of twp upgrades.
93#[derive(Debug, Clone)]
94pub struct InfoIterChain<A, B>(A, B);
95
96impl<A, B> Iterator for InfoIterChain<A, B>
97where
98    A: Iterator,
99    B: Iterator
100{
101    type Item = EitherName<A::Item, B::Item>;
102
103    fn next(&mut self) -> Option<Self::Item> {
104        if let Some(info) = self.0.next() {
105            return Some(EitherName::A(info))
106        }
107        if let Some(info) = self.1.next() {
108            return Some(EitherName::B(info))
109        }
110        None
111    }
112
113    fn size_hint(&self) -> (usize, Option<usize>) {
114        let (min1, max1) = self.0.size_hint();
115        let (min2, max2) = self.1.size_hint();
116        let max = max1.and_then(move |m1| max2.and_then(move |m2| m1.checked_add(m2)));
117        (min1.saturating_add(min2), max)
118    }
119}
120