fix: Split FlattenCompat logic to Flatten and FlatMap

Author: k-nasa

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

@@ -6,46 +6,6 @@ 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]
-        inner: FlattenCompat<Map<S, F, T, U>, 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 {
-            inner: FlattenCompat::new(stream.map(f)),
-        }
-    }
-}
-
-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>> {
-        self.project().inner.poll_next(cx)
-    }
-}
-
 pin_project! {
     /// This `struct` is created by the [`flatten`] method on [`Stream`]. See its
     /// documentation for more.
@@ -55,7 +15,9 @@ pin_project! {
     #[allow(missing_debug_implementations)]
     pub struct Flatten<S, U> {
         #[pin]
-        inner: FlattenCompat<S, U>
+        stream: S,
+        #[pin]
+        inner_stream: Option<U>,
     }
 }
 
@@ -66,47 +28,13 @@ where
 {
     pub(super) fn new(stream: S) -> Flatten<S, S::Item> {
         Flatten {
-            inner: FlattenCompat::new(stream),
-        }
-    }
-}
-
-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>> {
-        self.project().inner.poll_next(cx)
-    }
-}
-
-pin_project! {
-    /// Real logic of both `Flatten` and `FlatMap` which simply delegate to
-    /// this type.
-    #[derive(Clone, Debug)]
-    struct FlattenCompat<S, U> {
-        #[pin]
-        stream: S,
-        #[pin]
-        frontiter: Option<U>,
-    }
-}
-
-impl<S, U> FlattenCompat<S, U> {
-    /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`.
-    fn new(stream: S) -> FlattenCompat<S, U> {
-        FlattenCompat {
             stream,
-            frontiter: None,
+            inner_stream: None,
         }
     }
 }
 
-impl<S, U> Stream for FlattenCompat<S, U>
+impl<S, U> Stream for Flatten<S, <S::Item as IntoStream>::IntoStream>
 where
     S: Stream,
     S::Item: IntoStream<IntoStream = U, Item = U::Item>,
@@ -117,42 +45,16 @@ where
     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.frontiter.as_mut().as_pin_mut() {
+            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.frontiter.set(Some(inner.into_stream())),
+                Some(inner) => this.inner_stream.set(Some(inner.into_stream())),
             }
         }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::FlattenCompat;
-
-    use crate::prelude::*;
-    use crate::task;
-
-    use std::collections::VecDeque;
-
-    #[test]
-    fn test_poll_next() -> std::io::Result<()> {
-        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();
-
-        task::block_on(async move {
-            let flat = FlattenCompat::new(s);
-            let v: Vec<u8> = flat.collect().await;
-
-            assert_eq!(v, vec![1, 2, 3, 4, 5, 6]);
-            Ok(())
-        })
-    }
-}

src/stream/stream/mod.rs

@@ -99,10 +99,12 @@ cfg_unstable! {
     use crate::stream::into_stream::IntoStream;
 
     pub use merge::Merge;
-    pub use flatten::{FlatMap, Flatten};
+    pub use flatten::Flatten;
+    pub use flat_map::FlatMap;
 
     mod merge;
     mod flatten;
+    mod flat_map;
 }
 
 extension_trait! {