tokio_stream/stream_ext/
merge.rs

1use crate::stream_ext::Fuse;
2use crate::Stream;
3
4use core::pin::Pin;
5use core::task::{Context, Poll};
6use futures_core::FusedStream;
7use pin_project_lite::pin_project;
8
9pin_project! {
10    /// Stream returned by the [`merge`](super::StreamExt::merge) method.
11    pub struct Merge<T, U> {
12        #[pin]
13        a: Fuse<T>,
14        #[pin]
15        b: Fuse<U>,
16        // When `true`, poll `a` first, otherwise, `poll` b`.
17        a_first: bool,
18    }
19}
20
21impl<T, U> Merge<T, U> {
22    pub(super) fn new(a: T, b: U) -> Merge<T, U>
23    where
24        T: Stream,
25        U: Stream,
26    {
27        Merge {
28            a: Fuse::new(a),
29            b: Fuse::new(b),
30            a_first: true,
31        }
32    }
33}
34
35impl<T, U> Stream for Merge<T, U>
36where
37    T: Stream,
38    U: Stream<Item = T::Item>,
39{
40    type Item = T::Item;
41
42    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
43        let me = self.project();
44        let a_first = *me.a_first;
45
46        // Toggle the flag
47        *me.a_first = !a_first;
48
49        if a_first {
50            poll_next(me.a, me.b, cx)
51        } else {
52            poll_next(me.b, me.a, cx)
53        }
54    }
55
56    fn size_hint(&self) -> (usize, Option<usize>) {
57        super::merge_size_hints(self.a.size_hint(), self.b.size_hint())
58    }
59}
60
61impl<T, U> FusedStream for Merge<T, U>
62where
63    T: Stream,
64    U: Stream<Item = T::Item>,
65{
66    fn is_terminated(&self) -> bool {
67        self.a.is_terminated() && self.b.is_terminated()
68    }
69}
70
71fn poll_next<T, U>(
72    first: Pin<&mut T>,
73    second: Pin<&mut U>,
74    cx: &mut Context<'_>,
75) -> Poll<Option<T::Item>>
76where
77    T: Stream,
78    U: Stream<Item = T::Item>,
79{
80    let mut done = true;
81
82    match first.poll_next(cx) {
83        Poll::Ready(Some(val)) => return Poll::Ready(Some(val)),
84        Poll::Ready(None) => {}
85        Poll::Pending => done = false,
86    }
87
88    match second.poll_next(cx) {
89        Poll::Ready(Some(val)) => return Poll::Ready(Some(val)),
90        Poll::Ready(None) => {}
91        Poll::Pending => done = false,
92    }
93
94    if done {
95        Poll::Ready(None)
96    } else {
97        Poll::Pending
98    }
99}
tokio_stream/stream_ext/merge.rs

tokio_stream/stream_ext/
merge.rs
