tokio_stream/stream_ext/
then.rs

1use crate::Stream;
2
3use core::fmt;
4use core::future::Future;
5use core::pin::Pin;
6use core::task::{Context, Poll};
7use futures_core::FusedStream;
8use pin_project_lite::pin_project;
9
10pin_project! {
11    /// Stream for the [`then`](super::StreamExt::then) method.
12    #[must_use = "streams do nothing unless polled"]
13    pub struct Then<St, Fut, F> {
14        #[pin]
15        stream: St,
16        #[pin]
17        future: Option<Fut>,
18        f: F,
19    }
20}
21
22impl<St, Fut, F> fmt::Debug for Then<St, Fut, F>
23where
24    St: fmt::Debug,
25{
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("Then")
28            .field("stream", &self.stream)
29            .finish()
30    }
31}
32
33impl<St, Fut, F> Then<St, Fut, F> {
34    pub(super) fn new(stream: St, f: F) -> Self {
35        Then {
36            stream,
37            future: None,
38            f,
39        }
40    }
41}
42
43impl<St, F, Fut> Stream for Then<St, Fut, F>
44where
45    St: Stream,
46    Fut: Future,
47    F: FnMut(St::Item) -> Fut,
48{
49    type Item = Fut::Output;
50
51    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Fut::Output>> {
52        let mut me = self.project();
53
54        loop {
55            if let Some(future) = me.future.as_mut().as_pin_mut() {
56                match future.poll(cx) {
57                    Poll::Ready(item) => {
58                        me.future.set(None);
59                        return Poll::Ready(Some(item));
60                    }
61                    Poll::Pending => return Poll::Pending,
62                }
63            }
64
65            match me.stream.as_mut().poll_next(cx) {
66                Poll::Ready(Some(item)) => {
67                    me.future.set(Some((me.f)(item)));
68                }
69                Poll::Ready(None) => return Poll::Ready(None),
70                Poll::Pending => return Poll::Pending,
71            }
72        }
73    }
74
75    fn size_hint(&self) -> (usize, Option<usize>) {
76        let future_len = usize::from(self.future.is_some());
77        let (lower, upper) = self.stream.size_hint();
78
79        let lower = lower.saturating_add(future_len);
80        let upper = upper.and_then(|upper| upper.checked_add(future_len));
81
82        (lower, upper)
83    }
84}
85
86impl<St, F, Fut> FusedStream for Then<St, Fut, F>
87where
88    St: FusedStream,
89    Fut: Future,
90    F: FnMut(St::Item) -> Fut,
91{
92    fn is_terminated(&self) -> bool {
93        self.future.is_none() && self.stream.is_terminated()
94    }
95}
tokio_stream/stream_ext/then.rs

tokio_stream/stream_ext/
then.rs
