Add TryStreamExt::try_flatten (#1731)

* Add TryStreamExt::try_flatten

* Minor clean up StreamExt::flatten

Author: taiki-e

futures-util/src/stream/flatten.rs

@@ -9,26 +9,35 @@ use pin_utils::unsafe_pinned;
 #[derive(Debug)]
 #[must_use = "streams do nothing unless polled"]
 pub struct Flatten<St>
-    where St: Stream,
+where
+    St: Stream,
 {
     stream: St,
     next: Option<St::Item>,
 }
 
-impl<St: Stream> Unpin for Flatten<St>
-where St: Stream + Unpin,
-      St::Item: Stream + Unpin,
-{}
+impl<St> Unpin for Flatten<St>
+where
+    St: Stream + Unpin,
+    St::Item: Unpin,
+{
+}
 
-impl<St: Stream> Flatten<St>
-where St: Stream,
-      St::Item: Stream,
+impl<St> Flatten<St>
+where
+    St: Stream,
 {
     unsafe_pinned!(stream: St);
     unsafe_pinned!(next: Option<St::Item>);
+}
 
-    pub(super) fn new(stream: St) -> Flatten<St>{
-        Flatten { stream, next: None, }
+impl<St> Flatten<St>
+where
+    St: Stream,
+    St::Item: Stream,
+{
+    pub(super) fn new(stream: St) -> Self {
+        Self { stream, next: None }
     }
 
     /// Acquires a reference to the underlying stream that this combinator is
@@ -64,32 +73,33 @@ where St: Stream,
     }
 }
 
-impl<St: Stream + FusedStream> FusedStream for Flatten<St> {
+impl<St> FusedStream for Flatten<St>
+where
+    St: Stream + FusedStream,
+{
     fn is_terminated(&self) -> bool {
         self.next.is_none() && self.stream.is_terminated()
     }
 }
 
 impl<St> Stream for Flatten<St>
-    where St: Stream,
-          St::Item: Stream,
+where
+    St: Stream,
+    St::Item: Stream,
 {
     type Item = <St::Item as Stream>::Item;
 
-    fn poll_next(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Option<Self::Item>> {
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         loop {
             if self.next.is_none() {
                 match ready!(self.as_mut().stream().poll_next(cx)) {
                     Some(e) => self.as_mut().next().set(Some(e)),
                     None => return Poll::Ready(None),
                 }
             }
-            let item = ready!(self.as_mut().next().as_pin_mut().unwrap().poll_next(cx));
-            if item.is_some() {
-                return Poll::Ready(item);
+
+            if let Some(item) = ready!(self.as_mut().next().as_pin_mut().unwrap().poll_next(cx)) {
+                return Poll::Ready(Some(item));
             } else {
                 self.as_mut().next().set(None);
             }
@@ -100,8 +110,8 @@ impl<St> Stream for Flatten<St>
 // Forwarding impl of Sink from the underlying stream
 #[cfg(feature = "sink")]
 impl<S, Item> Sink<Item> for Flatten<S>
-    where S: Stream + Sink<Item>,
-          S::Item: Stream,
+where
+    S: Stream + Sink<Item>,
 {
     type Error = S::Error;
 

futures-util/src/try_stream/mod.rs

@@ -47,6 +47,9 @@ pub use self::try_filter::TryFilter;
 mod try_filter_map;
 pub use self::try_filter_map::TryFilterMap;
 
+mod try_flatten;
+pub use self::try_flatten::TryFlatten;
+
 mod try_collect;
 pub use self::try_collect::TryCollect;
 
@@ -556,6 +559,52 @@ pub trait TryStreamExt: TryStream {
         TryFilterMap::new(self, f)
     }
 
+    /// Flattens a stream of streams into just one continuous stream.
+    ///
+    /// If this stream's elements are themselves streams then this combinator
+    /// will flatten out the entire stream to one long chain of elements. Any
+    /// errors are passed through without looking at them, but otherwise each
+    /// individual stream will get exhausted before moving on to the next.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await)]
+    /// # futures::executor::block_on(async {
+    /// use futures::channel::mpsc;
+    /// use futures::stream::{StreamExt, TryStreamExt};
+    /// use std::thread;
+    ///
+    /// let (tx1, rx1) = mpsc::unbounded();
+    /// let (tx2, rx2) = mpsc::unbounded();
+    /// let (tx3, rx3) = mpsc::unbounded();
+    ///
+    /// thread::spawn(move || {
+    ///     tx1.unbounded_send(Ok(1)).unwrap();
+    /// });
+    /// thread::spawn(move || {
+    ///     tx2.unbounded_send(Ok(2)).unwrap();
+    ///     tx2.unbounded_send(Err(3)).unwrap();
+    /// });
+    /// thread::spawn(move || {
+    ///     tx3.unbounded_send(Ok(rx1)).unwrap();
+    ///     tx3.unbounded_send(Ok(rx2)).unwrap();
+    ///     tx3.unbounded_send(Err(4)).unwrap();
+    /// });
+    ///
+    /// let mut stream = rx3.try_flatten();
+    /// assert_eq!(stream.next().await, Some(Ok(1)));
+    /// assert_eq!(stream.next().await, Some(Ok(2)));
+    /// assert_eq!(stream.next().await, Some(Err(3)));
+    /// # });
+    /// ```
+    fn try_flatten(self) -> TryFlatten<Self>
+        where Self::Ok: TryStream,
+              <Self::Ok as TryStream>::Error: From<Self::Error>,
+              Self: Sized,
+    {
+        TryFlatten::new(self)
+    }
 
     /// Attempt to execute an accumulating asynchronous computation over a
     /// stream, collecting all the values into one final result.

futures-util/src/try_stream/try_flatten.rs

@@ -0,0 +1,127 @@
+use core::pin::Pin;
+use futures_core::stream::{FusedStream, Stream, TryStream};
+use futures_core::task::{Context, Poll};
+#[cfg(feature = "sink")]
+use futures_sink::Sink;
+use pin_utils::unsafe_pinned;
+
+/// Stream for the [`try_flatten`](super::TryStreamExt::try_flatten) method.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct TryFlatten<St>
+where
+    St: TryStream,
+{
+    stream: St,
+    next: Option<St::Ok>,
+}
+
+impl<St> Unpin for TryFlatten<St>
+where
+    St: TryStream + Unpin,
+    St::Ok: Unpin,
+{
+}
+
+impl<St> TryFlatten<St>
+where
+    St: TryStream,
+{
+    unsafe_pinned!(stream: St);
+    unsafe_pinned!(next: Option<St::Ok>);
+}
+
+impl<St> TryFlatten<St>
+where
+    St: TryStream,
+    St::Ok: TryStream,
+    <St::Ok as TryStream>::Error: From<St::Error>,
+{
+    pub(super) fn new(stream: St) -> Self {
+        Self { stream, next: None }
+    }
+
+    /// Acquires a reference to the underlying stream that this combinator is
+    /// pulling from.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Acquires a mutable reference to the underlying stream that this
+    /// combinator is pulling from.
+    ///
+    /// Note that care must be taken to avoid tampering with the state of the
+    /// stream which may otherwise confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Acquires a pinned mutable reference to the underlying stream that this
+    /// combinator is pulling from.
+    ///
+    /// Note that care must be taken to avoid tampering with the state of the
+    /// stream which may otherwise confuse this combinator.
+    pub fn get_pin_mut<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut St> {
+        self.stream()
+    }
+
+    /// Consumes this combinator, returning the underlying stream.
+    ///
+    /// Note that this may discard intermediate state of this combinator, so
+    /// care should be taken to avoid losing resources when this is called.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
+}
+
+impl<St> FusedStream for TryFlatten<St>
+where
+    St: TryStream + FusedStream,
+{
+    fn is_terminated(&self) -> bool {
+        self.next.is_none() && self.stream.is_terminated()
+    }
+}
+
+impl<St> Stream for TryFlatten<St>
+where
+    St: TryStream,
+    St::Ok: TryStream,
+    <St::Ok as TryStream>::Error: From<St::Error>,
+{
+    type Item = Result<<St::Ok as TryStream>::Ok, <St::Ok as TryStream>::Error>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            if self.next.is_none() {
+                match ready!(self.as_mut().stream().try_poll_next(cx)?) {
+                    Some(e) => self.as_mut().next().set(Some(e)),
+                    None => return Poll::Ready(None),
+                }
+            }
+
+            if let Some(item) = ready!(self
+                .as_mut()
+                .next()
+                .as_pin_mut()
+                .unwrap()
+                .try_poll_next(cx)?)
+            {
+                return Poll::Ready(Some(Ok(item)));
+            } else {
+                self.as_mut().next().set(None);
+            }
+        }
+    }
+}
+
+// Forwarding impl of Sink from the underlying stream
+#[cfg(feature = "sink")]
+impl<S, Item> Sink<Item> for TryFlatten<S>
+where
+    S: TryStream + Sink<Item>,
+{
+    type Error = <S as Sink<Item>>::Error;
+
+    delegate_sink!(stream, Item);
+}

futures/src/lib.rs

@@ -449,7 +449,7 @@ pub mod stream {
         TryStreamExt,
         AndThen, ErrInto, MapOk, MapErr, OrElse,
         InspectOk, InspectErr,
-        TryNext, TryForEach, TryFilterMap,
+        TryNext, TryForEach, TryFilterMap, TryFlatten,
         TryCollect, TryFold, TrySkipWhile,
         IntoStream,
     };