1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

use BufStream;

use either::Either;
use futures::Poll;

/// A buf stream that sequences two buf streams together.
///
/// `Chain` values are produced by the `chain` function on `BufStream`.
#[derive(Debug)]
pub struct Chain<T, U> {
    left: Option<T>,
    right: U,
}

impl<T, U> Chain<T, U> {
    pub(crate) fn new(left: T, right: U) -> Chain<T, U> {
        Chain {
            left: Some(left),
            right,
        }
    }
}

impl<T, U> BufStream for Chain<T, U>
where
    T: BufStream,
    U: BufStream<Error = T::Error>,
{
    type Item = Either<T::Item, U::Item>;
    type Error = T::Error;

    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        if let Some(ref mut stream) = self.left {
            let res = try_ready!(stream.poll_buf());

            if res.is_some() {
                return Ok(res.map(Either::Left).into());
            }
        }

        self.left = None;

        let res = try_ready!(self.right.poll_buf());
        Ok(res.map(Either::Right).into())
    }
}