implement concat() on fallible streams as TryStreamExt::try_concat()

Author: barzamin

futures-util/src/try_stream/mod.rs

@@ -31,6 +31,9 @@ pub use self::try_for_each::TryForEach;
 mod try_filter_map;
 pub use self::try_filter_map::TryFilterMap;
 
+mod try_concat;
+pub use self::try_concat::TryConcat;
+
 mod try_fold;
 pub use self::try_fold::TryFold;
 
@@ -438,6 +441,47 @@ pub trait TryStreamExt: TryStream {
         TryFold::new(self, f, init)
     }
 
+    /// Attempt to concatenate all items of a stream into a single
+    /// extendable destination, returning a future representing the end result.
+    ///
+    /// This combinator will extend the first item with the contents of all
+    /// the subsequent successful results of the stream. If the stream is empty,
+    /// the default value will be returned.
+    ///
+    /// Works with all collections that implement the [`Extend`](std::iter::Extend) trait.
+    ///
+    /// This method is similar to [`concat`](super::StreamExt::concat), but will
+    /// exit early if an error is encountered in the stream.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use futures::channel::mpsc;
+    /// use futures::executor::block_on;
+    /// use futures::stream::TryStreamExt;
+    /// use std::thread;
+    ///
+    /// let (mut tx, rx) = mpsc::unbounded::<Result<Vec<i32>, ()>>();
+    ///
+    /// thread::spawn(move || {
+    ///     for i in (0..3).rev() {
+    ///         let n = i * 3;
+    ///         tx.unbounded_send(Ok(vec![n + 1, n + 2, n + 3])).unwrap();
+    ///     }
+    /// });
+    ///
+    /// let result = block_on(rx.try_concat());
+    ///
+    /// assert_eq!(result, Ok(vec![7, 8, 9, 4, 5, 6, 1, 2, 3]));
+    /// ```
+    fn try_concat(self) -> TryConcat<Self>
+    where Self: Sized,
+          Self::Ok: Extend<<<Self as TryStream>::Ok as IntoIterator>::Item> +
+                    IntoIterator + Default,
+    {
+        TryConcat::new(self)
+    }
+
     /// Attempt to execute several futures from a stream concurrently.
     ///
     /// This stream's `Ok` type must be a [`TryFuture`] with an `Error` type

futures-util/src/try_stream/try_concat.rs

@@ -0,0 +1,63 @@
+use core::default::Default;
+use core::marker::Unpin;
+use core::pin::Pin;
+use futures_core::future::Future;
+use futures_core::stream::TryStream;
+use futures_core::task::{LocalWaker, Poll};
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+
+/// A stream combinator which attempts to concatenate the results of a stream into the
+/// first yielded item.
+///
+/// This structure is produced by the `TryStream::try_concat` method.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct TryConcat<St: TryStream> {
+    stream: St,
+    accum: Option<St::Ok>,
+}
+
+impl<St: TryStream + Unpin> Unpin for TryConcat<St> {}
+
+impl<St> TryConcat<St>
+where
+    St: TryStream,
+    St::Ok: Extend<<St::Ok as IntoIterator>::Item> + IntoIterator + Default,
+{
+    unsafe_pinned!(stream: St);
+    unsafe_unpinned!(accum: Option<St::Ok>);
+
+    pub(super) fn new(stream: St) -> TryConcat<St> {
+        TryConcat {
+            stream,
+            accum: None,
+        }
+    }
+}
+
+impl<St> Future for TryConcat<St>
+where
+    St: TryStream,
+    St::Ok: Extend<<St::Ok as IntoIterator>::Item> + IntoIterator + Default,
+{
+    type Output = Result<St::Ok, St::Error>;
+
+    fn poll(mut self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
+        loop {
+            match ready!(self.stream().try_poll_next(lw)) {
+                Some(Ok(x)) => {
+                    let accum = self.accum();
+                    if let Some(a) = accum {
+                        a.extend(x)
+                    } else {
+                        *accum = Some(x)
+                    }
+                },
+                Some(Err(e)) => return Poll::Ready(Err(e)),
+                None => {
+                    return Poll::Ready(Ok(self.accum().take().unwrap_or_default()))
+                }
+            }
+        }
+    }
+}