Implement stream combinators using async_stream_block and generators

Author: kpp

Cargo.toml

@@ -14,5 +14,8 @@ edition = "2018"
 pin-utils = "=0.1.0-alpha.4"
 
 [dependencies.futures]
-version = "=0.3.0-alpha.16"
+git = "https://github.com/taiki-e/futures-rs.git"
+branch = "async-stream"
+features = ["async-stream", "nightly"]
+#version = "=0.3.0-alpha.16"
 package = "futures-preview"

src/future.rs

@@ -1,5 +1,6 @@
 use futures::future::Future;
 use futures::stream::Stream;
+use futures::async_stream_block;
 
 use core::task::{Poll, Context};
 
@@ -100,35 +101,22 @@ 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
         }
-    })
+    }
 }
 
 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
+    }
 }
 
 pub fn poll_fn<F, T>(f: F) -> impl Future<Output = T>

src/lib.rs

@@ -1,4 +1,4 @@
-#![feature(async_await, 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 @@
 use futures::stream::Stream;
 use futures::future::Future;
+use futures::async_stream_block;
 
 use core::pin::Pin;
 use core::iter::IntoIterator;
@@ -30,48 +31,45 @@ 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)
+        }
+    }
 }
 
 pub fn filter<St, Fut, F>(stream: St, f: F) -> impl Stream<Item = St::Item>
     where St: Stream,
           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
-    })
+        }
+    }
 }
 
 pub fn filter_map<St, Fut, F, U>(stream: St, f: F) -> impl Stream<Item = U>
     where St: Stream,
           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
-    })
+        }
+    }
 }
 
 pub async fn into_future<St>(stream: St) -> (Option<St::Item>, impl Stream<Item = St::Item>)
@@ -121,18 +119,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
             }
         }
-    })
+    }
 }
 
 pub fn repeat<T>(item: T) -> impl Stream<Item = T>
@@ -148,147 +146,122 @@ 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;
         }
-    })
+    }
 }
 
 pub fn then<St, F, Fut>(stream: St, f: F) -> impl Stream<Item = St::Item>
     where St: Stream,
           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
         }
-    })
+    }
 }
 
 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
-        }
-    })
+    }
 }
 
 pub fn zip<St1, St2>(stream: St1, other: St2) -> impl Stream<Item = (St1::Item, St2::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,
+            }
         }
-    })
+    }
 }
 
 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
         }
-    })
+    }
 }
 
 pub fn take_while<St, F, Fut>(stream: St, f: F) -> impl Stream<Item = St::Item>
     where St: Stream,
           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
         }
-    })
+    }
 }
 
 pub fn skip_while<St, F, Fut>(stream: St, f: F) -> impl Stream<Item = St::Item>
     where St: Stream,
           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
-        }
-    })
+    }
 }
 
 pub async fn fold<St, T, F, Fut>(stream: St, init: T, f: F) -> T