Implement stream combinators using async_stream_block and generators

Author: Roman Proskuryakov

Cargo.toml

@@ -12,6 +12,7 @@ edition = "2018"
 
 [dependencies]
 pin-utils = "=0.1.0-alpha.4"
+futures-async-stream = "0.1.0-alpha.1"
 
 [dependencies.futures]
 version = "=0.3.0-alpha.17"

src/future.rs

@@ -1,5 +1,6 @@
 use futures::future::Future;
 use futures::stream::Stream;
+use futures_async_stream::async_stream_block;
 
 use core::task::{Poll, Context};
 
@@ -434,22 +435,13 @@ pub fn flatten_stream<Fut, St, T>(future: Fut) -> impl Stream<Item = T>
     where Fut: Future<Output = St>,
           St: Stream<Item = T>,
 {
-    use crate::stream::next;
-    futures::stream::unfold((Some(future), None), async move | (future, stream)| {
-        match (future, stream) {
-            (Some(future), None) => {
-                let stream = future.await;
-                let mut stream = Box::pin(stream);
-                let item = next(&mut stream).await;
-                item.map(|item| (item, (None, Some(stream))))
-            },
-            (None, Some(mut stream)) => {
-                let item = next(&mut stream).await;
-                item.map(|item| (item, (None, Some(stream))))
-            },
-            _ => unreachable!()
+    async_stream_block! {
+        let stream = future.await;
+        #[for_await]
+        for item in stream {
+            yield item
         }
-    })
+    }
 }
 
 /// Convert this future into a single element stream.
@@ -474,14 +466,10 @@ pub fn flatten_stream<Fut, St, T>(future: Fut) -> impl Stream<Item = T>
 pub fn into_stream<Fut>(future: Fut) -> impl Stream<Item = Fut::Output>
     where Fut: Future,
 {
-    futures::stream::unfold(Some(future), async move |future| {
-        if let Some(future) = future {
-            let item = future.await;
-            Some((item, (None)))
-        } else {
-            None
-        }
-    })
+    async_stream_block! {
+        let item = future.await;
+        yield item
+    }
 }
 
 /// Creates a new future wrapping around a function returning [`Poll`](core::task::Poll).

src/lib.rs

@@ -1,4 +1,4 @@
-#![feature(async_await, async_closure, gen_future, generators)]
+#![feature(async_await, gen_future, generators, proc_macro_hygiene)]
 
 pub mod future;
 pub mod stream;

src/stream.rs

@@ -1,5 +1,6 @@
 pub use futures::stream::Stream;
 use futures::future::Future;
+use futures_async_stream::async_stream_block;
 
 use core::pin::Pin;
 use core::iter::IntoIterator;
@@ -97,11 +98,13 @@ pub fn map<St, U, F>(stream: St, f: F) -> impl Stream<Item = U>
     where St: Stream,
           F: FnMut(St::Item) -> U,
 {
-    let stream = Box::pin(stream);
-    unfold((stream, f), async move | (mut stream, mut f)| {
-        let item = next(&mut stream).await;
-        item.map(|item| (f(item), (stream, f)))
-    })
+    let mut f = f;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
+            yield f(item)
+        }
+    }
 }
 
 /// Filters the values produced by this stream according to the provided
@@ -136,18 +139,15 @@ pub fn filter<St, Fut, F>(stream: St, f: F) -> impl Stream<Item = St::Item>
           F: FnMut(&St::Item) -> Fut,
           Fut: Future<Output = bool>
 {
-    let stream = Box::pin(stream);
-    unfold((stream, f), async move | (mut stream, mut f)| {
-        while let Some(item) = next(&mut stream).await {
-            let matched = f(&item).await;
-            if matched {
-                return Some((item, (stream, f)))
-            } else {
-                continue;
+    let mut f = f;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
+            if f(&item).await {
+                yield item
             }
-        };
-        None
-    })
+        }
+    }
 }
 
 /// Filters the values produced by this stream while simultaneously mapping
@@ -183,17 +183,15 @@ pub fn filter_map<St, Fut, F, U>(stream: St, f: F) -> impl Stream<Item = U>
           F: FnMut(St::Item) -> Fut,
           Fut: Future<Output = Option<U>>
 {
-    let stream = Box::pin(stream);
-    unfold((stream, f), async move | (mut stream, mut f)| {
-        while let Some(item) = next(&mut stream).await {
+    let mut f = f;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
             if let Some(item) = f(item).await {
-                return Some((item, (stream, f)))
-            } else {
-                continue;
+                yield item
             }
-        };
-        None
-    })
+        }
+    }
 }
 
 /// Converts this stream into a future of `(next_item, tail_of_stream)`.
@@ -370,18 +368,18 @@ pub async fn for_each<St, Fut, F>(stream: St, f: F) -> ()
 pub fn take<St>(stream: St, n: u64) -> impl Stream<Item = St::Item>
     where St: Stream,
 {
-    let stream = Box::pin(stream);
-    unfold((stream, n), async move | (mut stream, n)| {
-        if n == 0 {
-            None
-        } else {
-            if let Some(item) = next(&mut stream).await {
-                Some((item, (stream, n - 1)))
+    let mut n = n;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
+            if n == 0 {
+                break;
             } else {
-                None
+                n = n - 1;
+                yield item
             }
         }
-    })
+    }
 }
 
 /// Create a stream which produces the same item repeatedly.
@@ -435,27 +433,15 @@ pub fn flatten<St, SubSt, T>(stream: St) -> impl Stream<Item = T>
     where SubSt: Stream<Item = T>,
           St: Stream<Item = SubSt>,
 {
-    let stream = Box::pin(stream);
-    unfold((Some(stream), None), async move | (mut state_stream, mut state_substream)| {
-        loop {
-            if let Some(mut substream) = state_substream.take() {
-                if let Some(item) = next(&mut substream).await {
-                    return Some((item, (state_stream, Some(substream))))
-                } else {
-                    continue;
-                }
-            }
-            if let Some(mut stream) = state_stream.take() {
-                if let Some(substream) = next(&mut stream).await {
-                    let substream = Box::pin(substream);
-                    state_stream = Some(stream);
-                    state_substream = Some(substream);
-                    continue;
-                }
+    async_stream_block! {
+        #[for_await]
+        for substream in stream {
+            #[for_await]
+            for item in substream {
+                yield item
             }
-            return None;
         }
-    })
+    }
 }
 
 /// Computes from this stream's items new items of a different type using
@@ -488,16 +474,14 @@ pub fn then<St, F, Fut>(stream: St, f: F) -> impl Stream<Item = St::Item>
           F: FnMut(St::Item) -> Fut,
           Fut: Future<Output = St::Item>
 {
-    let stream = Box::pin(stream);
-    unfold((stream, f), async move | (mut stream, mut f)| {
-        let item = next(&mut stream).await;
-        if let Some(item) = item {
+    let mut f = f;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
             let new_item = f(item).await;
-            Some((new_item, (stream, f)))
-        } else {
-            None
+            yield new_item
         }
-    })
+    }
 }
 
 /// Creates a new stream which skips `n` items of the underlying stream.
@@ -522,22 +506,18 @@ pub fn then<St, F, Fut>(stream: St, f: F) -> impl Stream<Item = St::Item>
 pub fn skip<St>(stream: St, n: u64) -> impl Stream<Item = St::Item>
     where St: Stream,
 {
-    let stream = Box::pin(stream);
-    unfold((stream, n), async move | (mut stream, mut n)| {
-        while n != 0 {
-            if let Some(_) = next(&mut stream).await {
-                n = n - 1;
-                continue
+    let mut n = n;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
+            if n == 0 {
+                yield item
             } else {
-                return None
+                n = n - 1;
+                continue;
             }
         }
-        if let Some(item) = next(&mut stream).await {
-            Some((item, (stream, 0)))
-        } else {
-            None
-        }
-    })
+    }
 }
 
 /// An adapter for zipping two streams together.
@@ -566,16 +546,18 @@ pub fn zip<St1, St2>(stream: St1, other: St2) -> impl Stream<Item = (St1::Item,
     where St1: Stream,
           St2: Stream,
 {
-    let stream = Box::pin(stream);
-    let other = Box::pin(other);
-    unfold((stream, other), async move | (mut stream, mut other)| {
-        let left = next(&mut stream).await;
-        let right = next(&mut other).await;
-        match (left, right) {
-            (Some(left), Some(right)) => Some(((left, right), (stream, other))),
-            _ => None
+    let mut stream = Box::pin(stream);
+    let mut other = Box::pin(other);
+    async_stream_block! {
+        loop {
+            let left = next(&mut stream).await;
+            let right = next(&mut other).await;
+            match (left, right) {
+                (Some(left), Some(right)) => yield (left, right),
+                _ => break,
+            }
         }
-    })
+    }
 }
 
 /// Adapter for chaining two stream.
@@ -605,21 +587,16 @@ pub fn zip<St1, St2>(stream: St1, other: St2) -> impl Stream<Item = (St1::Item,
 pub fn chain<St>(stream: St, other: St) -> impl Stream<Item = St::Item>
     where St: Stream,
 {
-    let stream = Box::pin(stream);
-    let other = Box::pin(other);
-    let start_with_first = true;
-    unfold((stream, other, start_with_first), async move | (mut stream, mut other, start_with_first)| {
-        if start_with_first {
-            if let Some(item) = next(&mut stream).await {
-                return Some((item, (stream, other, start_with_first)))
-            }
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
+            yield item
         }
-        if let Some(item) = next(&mut other).await {
-            Some((item, (stream, other, /* start_with_first */ false)))
-        } else {
-            None
+        #[for_await]
+        for item in other {
+            yield item
         }
-    })
+    }
 }
 
 /// Take elements from this stream while the provided asynchronous predicate
@@ -649,18 +626,17 @@ pub fn take_while<St, F, Fut>(stream: St, f: F) -> impl Stream<Item = St::Item>
           F: FnMut(&St::Item) -> Fut,
           Fut: Future<Output = bool>,
 {
-    let stream = Box::pin(stream);
-    unfold((stream, f), async move | (mut stream, mut f)| {
-        if let Some(item) = next(&mut stream).await {
+    let mut f = f;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
             if f(&item).await {
-                Some((item, (stream, f)))
+                yield item
             } else {
-                None
+                break;
             }
-        } else {
-            None
         }
-    })
+    }
 }
 
 /// Skip elements on this stream while the provided asynchronous predicate
@@ -690,26 +666,23 @@ pub fn skip_while<St, F, Fut>(stream: St, f: F) -> impl Stream<Item = St::Item>
           F: FnMut(&St::Item) -> Fut,
           Fut: Future<Output = bool>,
 {
-    let stream = Box::pin(stream);
-    let should_skip = true;
-    unfold((stream, f, should_skip), async move | (mut stream, mut f, should_skip)| {
-        while should_skip {
-            if let Some(item) = next(&mut stream).await {
+    let mut f = f;
+    let mut should_skip = true;
+    async_stream_block! {
+        #[for_await]
+        for item in stream {
+            if should_skip {
                 if f(&item).await {
                     continue;
                 } else {
-                    return Some((item, (stream, f, /* should_skip */ false)))
+                    should_skip = false;
+                    yield item
                 }
             } else {
-                return None
+                yield item
             }
         }
-        if let Some(item) = next(&mut stream).await {
-            Some((item, (stream, f, /* should_skip */ false)))
-        } else {
-            None
-        }
-    })
+    }
 }
 
 /// Execute an accumulating asynchronous computation over a stream,