Merge #40

40: Add initial Fuse implementation for Stream r=yoshuawuyts a=spacejam

@matklad does this address your use case?

Co-authored-by: Tyler Neely <[email protected]>
Co-authored-by: Yoshua Wuyts <[email protected]>

Author: 3 people

src/stream/mod.rs

@@ -27,7 +27,7 @@ pub use from_stream::FromStream;
 pub use into_stream::IntoStream;
 pub use once::{once, Once};
 pub use repeat::{repeat, Repeat};
-pub use stream::{Scan, Stream, Take, Zip};
+pub use stream::{Fuse, Scan, Stream, Take, Zip};
 
 mod double_ended_stream;
 mod empty;

src/stream/stream/fuse.rs

@@ -0,0 +1,33 @@
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+/// A `Stream` that is permanently closed once a single call to `poll` results in
+/// `Poll::Ready(None)`, returning `Poll::Ready(None)` for all future calls to `poll`.
+#[derive(Clone, Debug)]
+pub struct Fuse<S> {
+    pub(crate) stream: S,
+    pub(crate) done: bool,
+}
+
+impl<S: Unpin> Unpin for Fuse<S> {}
+
+impl<S: futures_core::Stream> Fuse<S> {
+    pin_utils::unsafe_pinned!(stream: S);
+    pin_utils::unsafe_unpinned!(done: bool);
+}
+
+impl<S: futures_core::Stream> futures_core::Stream for Fuse<S> {
+    type Item = S::Item;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> {
+        if self.done {
+            Poll::Ready(None)
+        } else {
+            let next = futures_core::ready!(self.as_mut().stream().poll_next(cx));
+            if next.is_none() {
+                *self.as_mut().done() = true;
+            }
+            Poll::Ready(next)
+        }
+    }
+}

src/stream/stream/mod.rs

@@ -28,13 +28,15 @@ mod filter_map;
 mod find;
 mod find_map;
 mod fold;
+mod fuse;
 mod min_by;
 mod next;
 mod nth;
 mod scan;
 mod take;
 mod zip;
 
+pub use fuse::Fuse;
 pub use scan::Scan;
 pub use take::Take;
 pub use zip::Zip;
@@ -246,6 +248,35 @@ pub trait Stream {
         Enumerate::new(self)
     }
 
+    /// Transforms this `Stream` into a "fused" `Stream` such that after the first time `poll`
+    /// returns `Poll::Ready(None)`, all future calls to `poll` will also return
+    /// `Poll::Ready(None)`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # fn main() { async_std::task::block_on(async {
+    /// #
+    /// use async_std::prelude::*;
+    /// use async_std::stream;
+    ///
+    /// let mut s = stream::once(1).fuse();
+    /// assert_eq!(s.next().await, Some(1));
+    /// assert_eq!(s.next().await, None);
+    /// assert_eq!(s.next().await, None);
+    /// #
+    /// # }) }
+    /// ```
+    fn fuse(self) -> Fuse<Self>
+    where
+        Self: Sized,
+    {
+        Fuse {
+            stream: self,
+            done: false,
+        }
+    }
+
     /// Both filters and maps a stream.
     ///
     /// # Examples