Fixed `flatten`

src/stream/stream/flatten.rs

@@ -1,5 +1,5 @@
-use std::fmt;
+use core::fmt;
-use std::pin::Pin;
+use core::pin::Pin;
 
 use pin_project_lite::pin_project;
 
@@ -52,14 +52,21 @@ where
         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)) {
+                let next_item = futures_core::ready!(inner.poll_next(cx));
-                    return Poll::Ready(item);
+
+                if next_item.is_some() {
+                    return Poll::Ready(next_item);
+                } else {
+                    this.inner_stream.set(None);
                 }
             }
 
-            match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
+            let inner = 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())),
+            if inner.is_some() {
+                this.inner_stream.set(inner.map(IntoStream::into_stream));
+            } else {
+                return Poll::Ready(None);
             }
         }
     }

tests/stream.rs

@@ -1,3 +1,5 @@
+use std::convert::identity;
+use std::marker::Unpin;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
@@ -99,30 +101,21 @@ fn merge_works_with_unfused_streams() {
     assert_eq!(xs, vec![92, 92]);
 }
 
-#[test]
+struct S<T>(T);
-fn flat_map_doesnt_poll_completed_inner_stream() {
-    async_std::task::block_on(async {
-        use async_std::prelude::*;
-        use async_std::task::*;
-        use std::convert::identity;
-        use std::marker::Unpin;
-        use std::pin::Pin;
-
-        struct S<T>(T);
 
-        impl<T: Stream + Unpin> Stream for S<T> {
+impl<T: Stream + Unpin> Stream for S<T> {
     type Item = T::Item;
 
     fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
         unsafe { Pin::new_unchecked(&mut self.0) }.poll_next(ctx)
     }
-        }
+}
 
-        struct StrictOnce {
+struct StrictOnce {
     polled: bool,
-        };
+}
 
-        impl Stream for StrictOnce {
+impl Stream for StrictOnce {
     type Item = ();
 
     fn poll_next(mut self: Pin<&mut Self>, _: &mut Context) -> Poll<Option<Self::Item>> {
@@ -130,13 +123,13 @@ fn flat_map_doesnt_poll_completed_inner_stream() {
         self.polled = true;
         Poll::Ready(None)
     }
-        }
+}
 
-        struct Interchanger {
+struct Interchanger {
     polled: bool,
-        };
+}
 
-        impl Stream for Interchanger {
+impl Stream for Interchanger {
     type Item = S<Box<dyn Stream<Item = ()> + Unpin>>;
 
     fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
@@ -149,8 +142,11 @@ fn flat_map_doesnt_poll_completed_inner_stream() {
             Poll::Ready(Some(S(Box::new(StrictOnce { polled: false }))))
         }
     }
-        }
+}
 
+#[test]
+fn flat_map_doesnt_poll_completed_inner_stream() {
+    task::block_on(async {
         assert_eq!(
             Interchanger { polled: false }
                 .take(2)
@@ -161,3 +157,17 @@ fn flat_map_doesnt_poll_completed_inner_stream() {
         );
     });
 }
+
+#[test]
+fn flatten_doesnt_poll_completed_inner_stream() {
+    task::block_on(async {
+        assert_eq!(
+            Interchanger { polled: false }
+                .take(2)
+                .flatten()
+                .count()
+                .await,
+            0
+        );
+    });
+}