From bb1416420d047638d2cc8be21920206c3053870a Mon Sep 17 00:00:00 2001
From: k-nasa <htilcs1115@gmail.com>
Date: Thu, 17 Oct 2019 23:56:32 +0900
Subject: [PATCH] feat: Add Stream trait for FlattenCompat

---
 src/stream/stream/flatten.rs | 49 ++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/src/stream/stream/flatten.rs b/src/stream/stream/flatten.rs
index fffda4d..7265d17 100644
--- a/src/stream/stream/flatten.rs
+++ b/src/stream/stream/flatten.rs
@@ -22,3 +22,52 @@ impl<S, U> FlattenCompat<S, U> {
         }
     }
 }
+
+impl<S, U> Stream for FlattenCompat<S, U>
+where
+    S: Stream<Item: IntoStream<IntoStream = U, Item = U::Item>> + std::marker::Unpin,
+    U: Stream + std::marker::Unpin,
+{
+    type Item = U::Item;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            if let Some(ref mut inner) = self.as_mut().frontiter() {
+                if let item @ Some(_) = futures_core::ready!(Pin::new(inner).poll_next(cx)) {
+                    return Poll::Ready(item);
+                }
+            }
+
+            match futures_core::ready!(Pin::new(&mut self.stream).poll_next(cx)) {
+                None => return Poll::Ready(None),
+                Some(inner) => *self.as_mut().frontiter() = 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(())
+        })
+    }
+}