rocket/util/
join.rs

Help
1use std::pin::Pin;
2use std::task::{Poll, Context};
3
4use pin_project_lite::pin_project;
5
6use futures::stream::Stream;
7use futures::ready;
8
9/// Join two streams, `a` and `b`, into a new `Join` stream that returns items
10/// from both streams, biased to `a`, until `a` finishes. The joined stream
11/// completes when `a` completes, irrespective of `b`. If `b` stops producing
12/// values, then the joined stream acts exactly like a fused `a`.
13///
14/// Values are biased to those of `a`: if `a` provides a value, it is always
15/// emitted before a value provided by `b`. In other words, values from `b` are
16/// emitted when and only when `a` is not producing a value.
17pub fn join<A: Stream, B: Stream>(a: A, b: B) -> Join<A, B> {
18    Join { a, b: Some(b), done: false, }
19}
20
21pin_project! {
22    /// Stream returned by [`join`].
23    pub struct Join<T, U> {
24        #[pin]
25        a: T,
26        #[pin]
27        b: Option<U>,
28        // Set when `a` returns `None`.
29        done: bool,
30    }
31}
32
33impl<T, U> Stream for Join<T, U>
34    where T: Stream,
35          U: Stream<Item = T::Item>,
36{
37    type Item = T::Item;
38
39    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
40        if self.done {
41            return Poll::Ready(None);
42        }
43
44        let me = self.as_mut().project();
45        match me.a.poll_next(cx) {
46            Poll::Ready(opt) => {
47                *me.done = opt.is_none();
48                Poll::Ready(opt)
49            },
50            Poll::Pending => match me.b.as_pin_mut() {
51                None => Poll::Pending,
52                Some(b) => match ready!(b.poll_next(cx)) {
53                    Some(value) => Poll::Ready(Some(value)),
54                    None => {
55                        self.as_mut().project().b.set(None);
56                        Poll::Pending
57                    }
58                }
59            }
60        }
61    }
62
63    fn size_hint(&self) -> (usize, Option<usize>) {
64        let (left_low, left_high) = self.a.size_hint();
65        let (right_low, right_high) = self.b.as_ref()
66            .map(|b| b.size_hint())
67            .unwrap_or_default();
68
69        let low = left_low.saturating_add(right_low);
70        let high = match (left_high, right_high) {
71            (Some(h1), Some(h2)) => h1.checked_add(h2),
72            _ => None,
73        };
74
75        (low, high)
76    }
77}
rocket/util/join.rs

rocket/util/
join.rs
