Merge pull request #73 from yoshuawuyts/zip-and-chain-traits

Init zip and chain traits

Author: yoshuawuyts

src/lib.rs

@@ -57,8 +57,10 @@ pub mod prelude {
     pub use super::future::Race as _;
     pub use super::future::RaceOk as _;
     pub use super::future::TryJoin as _;
+    pub use super::stream::Chain as _;
     pub use super::stream::IntoStream as _;
     pub use super::stream::Merge as _;
+    pub use super::stream::Zip as _;
 }
 
 pub mod future;
@@ -70,13 +72,17 @@ pub mod array {
     pub use crate::future::race::array::Race;
     pub use crate::future::race_ok::array::{AggregateError, RaceOk};
     pub use crate::future::try_join::array::TryJoin;
+    pub use crate::stream::chain::array::Chain;
     pub use crate::stream::merge::array::Merge;
+    pub use crate::stream::zip::array::Zip;
 }
 /// A contiguous growable array type with heap-allocated contents, written `Vec<T>`.
 pub mod vec {
     pub use crate::future::join::vec::Join;
     pub use crate::future::race::vec::Race;
     pub use crate::future::race_ok::vec::{AggregateError, RaceOk};
     pub use crate::future::try_join::vec::TryJoin;
+    pub use crate::stream::chain::vec::Chain;
     pub use crate::stream::merge::vec::Merge;
+    pub use crate::stream::zip::vec::Zip;
 }

src/stream/chain/array.rs

@@ -0,0 +1,101 @@
+use core::fmt;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+
+use futures_core::Stream;
+use pin_project::pin_project;
+
+use crate::utils;
+
+use super::Chain as ChainTrait;
+
+/// A stream that chains multiple streams one after another.
+///
+/// This `struct` is created by the [`chain`] method on the [`Chain`] trait. See its
+/// documentation for more.
+///
+/// [`chain`]: trait.Chain.html#method.merge
+/// [`Chain`]: trait.Chain.html
+#[pin_project]
+pub struct Chain<S, const N: usize> {
+    #[pin]
+    streams: [S; N],
+    index: usize,
+    len: usize,
+    done: bool,
+}
+
+impl<S: Stream, const N: usize> Stream for Chain<S, N> {
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+
+        assert!(!*this.done, "Stream should not be polled after completion");
+
+        loop {
+            if this.index == this.len {
+                *this.done = true;
+                return Poll::Ready(None);
+            }
+            let stream = utils::iter_pin_mut(this.streams.as_mut())
+                .nth(*this.index)
+                .unwrap();
+            match stream.poll_next(cx) {
+                Poll::Ready(Some(item)) => return Poll::Ready(Some(item)),
+                Poll::Ready(None) => {
+                    *this.index += 1;
+                    continue;
+                }
+                Poll::Pending => return Poll::Pending,
+            }
+        }
+    }
+}
+
+impl<S, const N: usize> fmt::Debug for Chain<S, N>
+where
+    S: Stream + fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_list().entries(self.streams.iter()).finish()
+    }
+}
+
+impl<S: Stream, const N: usize> ChainTrait for [S; N] {
+    type Item = S::Item;
+
+    type Stream = Chain<S, N>;
+
+    fn chain(self) -> Self::Stream {
+        Chain {
+            len: self.len(),
+            streams: self,
+            index: 0,
+            done: false,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures_lite::future::block_on;
+    use futures_lite::prelude::*;
+    use futures_lite::stream;
+
+    #[test]
+    fn chain_3() {
+        block_on(async {
+            let a = stream::once(1);
+            let b = stream::once(2);
+            let c = stream::once(3);
+            let mut s = [a, b, c].chain();
+
+            assert_eq!(s.next().await, Some(1));
+            assert_eq!(s.next().await, Some(2));
+            assert_eq!(s.next().await, Some(3));
+            assert_eq!(s.next().await, None);
+        })
+    }
+}

src/stream/chain/mod.rs

@@ -0,0 +1,17 @@
+use futures_core::Stream;
+
+pub(crate) mod array;
+pub(crate) mod tuple;
+pub(crate) mod vec;
+
+/// Takes multiple streams and creates a new stream over all in sequence.
+pub trait Chain {
+    /// What's the return type of our stream?
+    type Item;
+
+    /// What stream do we return?
+    type Stream: Stream<Item = Self::Item>;
+
+    /// Combine multiple streams into a single stream.
+    fn chain(self) -> Self::Stream;
+}

src/stream/chain/tuple.rs

@@ -0,0 +1 @@
+

src/stream/chain/vec.rs

@@ -0,0 +1,101 @@
+use core::fmt;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+
+use futures_core::Stream;
+use pin_project::pin_project;
+
+use crate::utils;
+
+use super::Chain as ChainTrait;
+
+/// A stream that chains multiple streams one after another.
+///
+/// This `struct` is created by the [`chain`] method on the [`Chain`] trait. See its
+/// documentation for more.
+///
+/// [`chain`]: trait.Chain.html#method.merge
+/// [`Chain`]: trait.Chain.html
+#[pin_project]
+pub struct Chain<S> {
+    #[pin]
+    streams: Vec<S>,
+    index: usize,
+    len: usize,
+    done: bool,
+}
+
+impl<S: Stream> Stream for Chain<S> {
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+
+        assert!(!*this.done, "Stream should not be polled after completion");
+
+        loop {
+            if this.index == this.len {
+                *this.done = true;
+                return Poll::Ready(None);
+            }
+            let stream = utils::iter_pin_mut_vec(this.streams.as_mut())
+                .nth(*this.index)
+                .unwrap();
+            match stream.poll_next(cx) {
+                Poll::Ready(Some(item)) => return Poll::Ready(Some(item)),
+                Poll::Ready(None) => {
+                    *this.index += 1;
+                    continue;
+                }
+                Poll::Pending => return Poll::Pending,
+            }
+        }
+    }
+}
+
+impl<S> fmt::Debug for Chain<S>
+where
+    S: Stream + fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_list().entries(self.streams.iter()).finish()
+    }
+}
+
+impl<S: Stream> ChainTrait for Vec<S> {
+    type Item = S::Item;
+
+    type Stream = Chain<S>;
+
+    fn chain(self) -> Self::Stream {
+        Chain {
+            len: self.len(),
+            streams: self,
+            index: 0,
+            done: false,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures_lite::future::block_on;
+    use futures_lite::prelude::*;
+    use futures_lite::stream;
+
+    #[test]
+    fn chain_3() {
+        block_on(async {
+            let a = stream::once(1);
+            let b = stream::once(2);
+            let c = stream::once(3);
+            let mut s = vec![a, b, c].chain();
+
+            assert_eq!(s.next().await, Some(1));
+            assert_eq!(s.next().await, Some(2));
+            assert_eq!(s.next().await, Some(3));
+            assert_eq!(s.next().await, None);
+        })
+    }
+}

src/stream/mod.rs

@@ -47,8 +47,12 @@
 //!
 //! See the [future concurrency][crate::future#concurrency] documentation for
 //! more on futures concurrency.
+pub use chain::Chain;
 pub use into_stream::IntoStream;
 pub use merge::Merge;
+pub use zip::Zip;
 
+pub(crate) mod chain;
 mod into_stream;
 pub(crate) mod merge;
+pub(crate) mod zip;