tokio_stream/stream_ext/
fuse.rs

1use crate::Stream;
2
3use futures_core::FusedStream;
4use pin_project_lite::pin_project;
5use std::pin::Pin;
6use std::task::{ready, Context, Poll};
7
8pin_project! {
9    /// Stream returned by [`fuse()`][super::StreamExt::fuse].
10    #[derive(Debug)]
11    pub struct Fuse<T> {
12        #[pin]
13        stream: Option<T>,
14    }
15}
16
17impl<T> Fuse<T>
18where
19    T: Stream,
20{
21    pub(crate) fn new(stream: T) -> Fuse<T> {
22        Fuse {
23            stream: Some(stream),
24        }
25    }
26}
27
28impl<T> Stream for Fuse<T>
29where
30    T: Stream,
31{
32    type Item = T::Item;
33
34    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
35        let res = match Option::as_pin_mut(self.as_mut().project().stream) {
36            Some(stream) => ready!(stream.poll_next(cx)),
37            None => return Poll::Ready(None),
38        };
39
40        if res.is_none() {
41            // Do not poll the stream anymore
42            self.as_mut().project().stream.set(None);
43        }
44
45        Poll::Ready(res)
46    }
47
48    fn size_hint(&self) -> (usize, Option<usize>) {
49        match self.stream {
50            Some(ref stream) => stream.size_hint(),
51            None => (0, Some(0)),
52        }
53    }
54}
55
56impl<T> FusedStream for Fuse<T>
57where
58    T: Stream,
59{
60    fn is_terminated(&self) -> bool {
61        self.stream.is_none()
62    }
63}
tokio_stream/stream_ext/fuse.rs

tokio_stream/stream_ext/
fuse.rs
