futures_util/stream/
zip.rs

1use futures_core::{Async, Poll, Stream};
2use futures_core::task;
3
4use stream::{StreamExt, Fuse};
5
6/// An adapter for merging the output of two streams.
7///
8/// The merged stream produces items from one or both of the underlying
9/// streams as they become available. Errors, however, are not merged: you
10#[derive(Debug)]
11/// get at most one error at a time.
12#[must_use = "streams do nothing unless polled"]
13pub struct Zip<S1: Stream, S2: Stream> {
14    stream1: Fuse<S1>,
15    stream2: Fuse<S2>,
16    queued1: Option<S1::Item>,
17    queued2: Option<S2::Item>,
18}
19
20pub fn new<S1, S2>(stream1: S1, stream2: S2) -> Zip<S1, S2>
21    where S1: Stream, S2: Stream<Error = S1::Error>
22{
23    Zip {
24        stream1: stream1.fuse(),
25        stream2: stream2.fuse(),
26        queued1: None,
27        queued2: None,
28    }
29}
30
31impl<S1, S2> Stream for Zip<S1, S2>
32    where S1: Stream, S2: Stream<Error = S1::Error>
33{
34    type Item = (S1::Item, S2::Item);
35    type Error = S1::Error;
36
37    fn poll_next(&mut self, cx: &mut task::Context) -> Poll<Option<Self::Item>, Self::Error> {
38        if self.queued1.is_none() {
39            match self.stream1.poll_next(cx)? {
40                Async::Ready(Some(item1)) => self.queued1 = Some(item1),
41                Async::Ready(None) | Async::Pending => {}
42            }
43        }
44        if self.queued2.is_none() {
45            match self.stream2.poll_next(cx)? {
46                Async::Ready(Some(item2)) => self.queued2 = Some(item2),
47                Async::Ready(None) | Async::Pending => {}
48            }
49        }
50
51        if self.queued1.is_some() && self.queued2.is_some() {
52            let pair = (self.queued1.take().unwrap(),
53                        self.queued2.take().unwrap());
54            Ok(Async::Ready(Some(pair)))
55        } else if self.stream1.is_done() || self.stream2.is_done() {
56            Ok(Async::Ready(None))
57        } else {
58            Ok(Async::Pending)
59        }
60    }
61}