Merge pull request #367 from k-nasa/add_stream_flatten

Add Stream::flatten and Stream::flat_map

Author: yoshuawuyts

src/stream/stream/flat_map.rs

@@ -0,0 +1,62 @@
+use pin_project_lite::pin_project;
+use std::pin::Pin;
+
+use crate::prelude::*;
+use crate::stream::stream::map::Map;
+use crate::stream::{IntoStream, Stream};
+use crate::task::{Context, Poll};
+
+pin_project! {
+    /// This `struct` is created by the [`flat_map`] method on [`Stream`]. See its
+    /// documentation for more.
+    ///
+    /// [`flat_map`]: trait.Stream.html#method.flat_map
+    /// [`Stream`]: trait.Stream.html
+    #[allow(missing_debug_implementations)]
+    pub struct FlatMap<S, U, T, F> {
+        #[pin]
+        stream: Map<S, F, T, U>,
+        #[pin]
+        inner_stream: Option<U>,
+    }
+}
+
+impl<S, U, F> FlatMap<S, U, S::Item, F>
+where
+    S: Stream,
+    U: IntoStream,
+    F: FnMut(S::Item) -> U,
+{
+    pub(super) fn new(stream: S, f: F) -> FlatMap<S, U, S::Item, F> {
+        FlatMap {
+            stream: stream.map(f),
+            inner_stream: None,
+        }
+    }
+}
+
+impl<S, U, F> Stream for FlatMap<S, U, S::Item, F>
+where
+    S: Stream,
+    S::Item: IntoStream<IntoStream = U, Item = U::Item>,
+    U: Stream,
+    F: FnMut(S::Item) -> U,
+{
+    type Item = U::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+        loop {
+            if let Some(inner) = this.inner_stream.as_mut().as_pin_mut() {
+                if let item @ Some(_) = futures_core::ready!(inner.poll_next(cx)) {
+                    return Poll::Ready(item);
+                }
+            }
+
+            match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
+                None => return Poll::Ready(None),
+                Some(inner) => this.inner_stream.set(Some(inner.into_stream())),
+            }
+        }
+    }
+}

src/stream/stream/flatten.rs

@@ -0,0 +1,58 @@
+use pin_project_lite::pin_project;
+use std::pin::Pin;
+
+use crate::stream::{IntoStream, Stream};
+use crate::task::{Context, Poll};
+
+pin_project! {
+    /// This `struct` is created by the [`flatten`] method on [`Stream`]. See its
+    /// documentation for more.
+    ///
+    /// [`flatten`]: trait.Stream.html#method.flatten
+    /// [`Stream`]: trait.Stream.html
+    #[allow(missing_debug_implementations)]
+    pub struct Flatten<S, U> {
+        #[pin]
+        stream: S,
+        #[pin]
+        inner_stream: Option<U>,
+    }
+}
+
+impl<S> Flatten<S, S::Item>
+where
+    S: Stream,
+    S::Item: IntoStream,
+{
+    pub(super) fn new(stream: S) -> Flatten<S, S::Item> {
+        Flatten {
+            stream,
+            inner_stream: None,
+        }
+    }
+}
+
+impl<S, U> Stream for Flatten<S, <S::Item as IntoStream>::IntoStream>
+where
+    S: Stream,
+    S::Item: IntoStream<IntoStream = U, Item = U::Item>,
+    U: Stream,
+{
+    type Item = U::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+        loop {
+            if let Some(inner) = this.inner_stream.as_mut().as_pin_mut() {
+                if let item @ Some(_) = futures_core::ready!(inner.poll_next(cx)) {
+                    return Poll::Ready(item);
+                }
+            }
+
+            match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
+                None => return Poll::Ready(None),
+                Some(inner) => this.inner_stream.set(Some(inner.into_stream())),
+            }
+        }
+    }
+}

src/stream/stream/mod.rs

@@ -96,13 +96,17 @@ cfg_unstable! {
     use std::time::Duration;
 
     use crate::future::Future;
-    use crate::stream::FromStream;
-    use crate::stream::{Product, Sum};
+    use crate::stream::into_stream::IntoStream;
+    use crate::stream::{FromStream, Product, Sum};
 
     pub use merge::Merge;
+    pub use flatten::Flatten;
+    pub use flat_map::FlatMap;
     pub use timeout::{TimeoutError, Timeout};
 
     mod merge;
+    mod flatten;
+    mod flat_map;
     mod timeout;
 }
 
@@ -563,6 +567,76 @@ extension_trait! {
             Filter::new(self, predicate)
         }
 
+        #[doc= r#"
+            Creates an stream that works like map, but flattens nested structure.
+
+            # Examples
+
+            Basic usage:
+
+            ```
+            # async_std::task::block_on(async {
+
+            use std::collections::VecDeque;
+            use async_std::prelude::*;
+            use async_std::stream::IntoStream;
+
+            let inner1: VecDeque<u8> = vec![1,2,3].into_iter().collect();
+            let inner2: VecDeque<u8> = vec![4,5,6].into_iter().collect();
+
+            let s: VecDeque<_> = vec![inner1, inner2].into_iter().collect();
+
+            let v :Vec<_> = s.flat_map(|s| s.into_stream()).collect().await;
+
+            assert_eq!(v, vec![1,2,3,4,5,6]);
+
+            # });
+            ```
+        "#]
+        #[cfg(feature = "unstable")]
+        #[cfg_attr(feature = "docs", doc(cfg(unstable)))]
+        fn flat_map<U, F>(self, f: F) -> FlatMap<Self, U, Self::Item, F>
+            where
+                Self: Sized,
+                U: IntoStream,
+                F: FnMut(Self::Item) -> U,
+        {
+            FlatMap::new(self, f)
+        }
+
+        #[doc = r#"
+            Creates an stream that flattens nested structure.
+
+            # Examples
+
+            Basic usage:
+
+            ```
+            # async_std::task::block_on(async {
+
+            use std::collections::VecDeque;
+            use async_std::prelude::*;
+
+            let inner1: VecDeque<u8> = vec![1,2,3].into_iter().collect();
+            let inner2: VecDeque<u8> = vec![4,5,6].into_iter().collect();
+            let s: VecDeque<_> = vec![inner1, inner2].into_iter().collect();
+
+            let v: Vec<_> = s.flatten().collect().await;
+
+            assert_eq!(v, vec![1,2,3,4,5,6]);
+
+            # });
+        "#]
+        #[cfg(feature = "unstable")]
+        #[cfg_attr(feature = "docs", doc(cfg(unstable)))]
+        fn flatten(self) -> Flatten<Self, Self::Item>
+        where
+            Self: Sized,
+            Self::Item: IntoStream,
+        {
+            Flatten::new(self)
+        }
+
         #[doc = r#"
             Both filters and maps a stream.