diff --git a/src/stream/stream/mod.rs b/src/stream/stream/mod.rs index d140228..e68e6ac 100644 --- a/src/stream/stream/mod.rs +++ b/src/stream/stream/mod.rs @@ -123,6 +123,7 @@ cfg_unstable! { use count::CountFuture; use partition::PartitionFuture; + use unzip::UnzipFuture; pub use merge::Merge; pub use flatten::Flatten; @@ -137,6 +138,7 @@ cfg_unstable! { mod partition; mod timeout; mod throttle; + mod unzip; } extension_trait! { @@ -1748,6 +1750,44 @@ extension_trait! { Zip::new(self, other) } + #[doc = r#" + Converts an stream of pairs into a pair of containers. + + `unzip()` consumes an entire stream of pairs, producing two collections: one from the left elements of the pairs, and one from the right elements. + + This function is, in some sense, the opposite of [`zip`]. + + [`zip`]: trait.Stream.html#method.zip + + # Example + + ``` + # fn main() { async_std::task::block_on(async { + # + use async_std::prelude::*; + use async_std::stream; + + let s = stream::from_iter(vec![(1,2), (3,4)]); + + let (left, right): (Vec<_>, Vec<_>) = s.unzip().await; + + assert_eq!(left, [1, 3]); + assert_eq!(right, [2, 4]); + # + # }) } + ``` + "#] + #[cfg(feature = "unstable")] + #[cfg_attr(feature = "docs", doc(cfg(unstable)))] + fn unzip(self) -> impl Future [UnzipFuture] + where + FromA: Default + Extend, + FromB: Default + Extend, + Self: Stream + Sized, + { + UnzipFuture::new(self) + } + #[doc = r#" Transforms a stream into a collection. diff --git a/src/stream/stream/unzip.rs b/src/stream/stream/unzip.rs new file mode 100644 index 0000000..e0832ff --- /dev/null +++ b/src/stream/stream/unzip.rs @@ -0,0 +1,57 @@ +use std::future::Future; +use std::pin::Pin; + +use pin_project_lite::pin_project; + +use crate::stream::Stream; +use crate::task::{Context, Poll}; + +pin_project! { + #[derive(Clone, Debug)] + #[cfg(all(feature = "default", feature = "unstable"))] + #[cfg_attr(feature = "docs", doc(cfg(unstable)))] + pub struct UnzipFuture { + #[pin] + stream: S, + res: Option<(FromA, FromB)>, + } +} + +impl UnzipFuture +where + FromA: Default, + FromB: Default, +{ + pub(super) fn new(stream: S) -> Self { + UnzipFuture { + stream, + res: Some((FromA::default(), FromB::default())), + } + } +} + +impl Future for UnzipFuture +where + S: Stream, + FromA: Default + Extend, + FromB: Default + Extend, +{ + type Output = (FromA, FromB); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut this = self.project(); + + loop { + let next = futures_core::ready!(this.stream.as_mut().poll_next(cx)); + + match next { + Some((a, b)) => { + let res = this.res.as_mut().unwrap(); + res.0.extend(Some(a)); + res.1.extend(Some(b)); + } + None => return Poll::Ready(this.res.take().unwrap()), + } + } + } +}