Repository: alexcrichton/futures-rs
Pull request #545 (2017-07-31 07:40:37)

Implement stream::poll_fn
WIP. Using CI because I'm affected by https://github.com/rust-lang/rust/issues/43132.

EDIT: seems passed.

alexcrichton: Shouldn't this be included in `stream/mod.rs` as well? Can you be sure to test this patch locally and include a unit test as well?
ishitatsuyuki: Uh, can't right now. I will have to do a rebuild.

I modified mod.rs initially but I accidentally removed it alongside with whitespace change.

I will address the issues including tests tomorrow.

Patch:
diff --git a/src/stream/mod.rs b/src/stream/mod.rs
index cbe09e4d..ee0dfaf7 100755
--- a/src/stream/mod.rs
+++ b/src/stream/mod.rs
@@ -43,6 +43,7 @@ mod merge;
 mod once;
 mod or_else;
 mod peek;
+mod poll_fn;
 mod select;
 mod skip;
 mod skip_while;
@@ -70,6 +71,7 @@ pub use self::merge::{Merge, MergedItem};
 pub use self::once::{Once, once};
 pub use self::or_else::OrElse;
 pub use self::peek::Peekable;
+pub use self::poll_fn::{poll_fn, PollFn};
 pub use self::select::Select;
 pub use self::skip::Skip;
 pub use self::skip_while::SkipWhile;
diff --git a/src/stream/poll_fn.rs b/src/stream/poll_fn.rs
new file mode 100644
index 00000000..fbc7df08
--- /dev/null
+++ b/src/stream/poll_fn.rs
@@ -0,0 +1,49 @@
+//! Definition of the `PollFn` combinator
+
+use {Stream, Poll};
+
+/// A stream which adapts a function returning `Poll`.
+///
+/// Created by the `poll_fn` function.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct PollFn<F> {
+    inner: F,
+}
+
+/// Creates a new stream wrapping around a function returning `Poll`.
+///
+/// Polling the returned stream delegates to the wrapped function.
+///
+/// # Examples
+///
+/// ```
+/// use futures::stream::poll_fn;
+/// use futures::{Async, Poll};
+///
+/// let mut counter = 1usize;
+///
+/// let read_stream = poll_fn(move || -> Poll<Option<String>, std::io::Error> {
+///     if counter == 0 { return Ok(Async::Ready(None)); }
+///     counter -= 1;
+///     Ok(Async::Ready(Some("Hello, World!".to_owned())))
+/// });
+/// ```
+pub fn poll_fn<T, E, F>(f: F) -> PollFn<F>
+where
+    F: FnMut() -> Poll<Option<T>, E>,
+{
+    PollFn { inner: f }
+}
+
+impl<T, E, F> Stream for PollFn<F>
+where
+    F: FnMut() -> Poll<Option<T>, E>,
+{
+    type Item = T;
+    type Error = E;
+
+    fn poll(&mut self) -> Poll<Option<T>, E> {
+        (self.inner)()
+    }
+}
diff --git a/src/sync/mpsc/mod.rs b/src/sync/mpsc/mod.rs
index b7e1eadc..670d7384 100644
--- a/src/sync/mpsc/mod.rs
+++ b/src/sync/mpsc/mod.rs
@@ -78,9 +78,7 @@ use std::usize;
 
 use sync::mpsc::queue::{Queue, PopResult};
 use task::{self, Task};
-use future::Executor;
-use sink::SendAll;
-use {Async, AsyncSink, Future, Poll, StartSend, Sink, Stream};
+use {Async, AsyncSink, Poll, StartSend, Sink, Stream};
 
 mod queue;
 
@@ -826,141 +824,6 @@ impl<T> Stream for UnboundedReceiver<T> {
     }
 }
 
-/// Handle returned from the `spawn` function.
-///
-/// This handle is a stream that proxies a stream on a seperate `Executor`.
-/// Created through the `mpsc::spawn` function, this handle will produce
-/// the same values as the proxied stream, as they are produced in the executor,
-/// and uses a limited buffer to exert back-pressure on the remote stream.
-///
-/// If this handle is dropped, then the stream will no longer be polled and is
-/// scheduled to be dropped.
-pub struct SpawnHandle<Item, Error> {
-    rx: Receiver<Result<Item, Error>>
-}
-
-/// Type of future which `Executor` instances must be able to execute for `spawn`.
-pub struct Execute<S: Stream> {
-    inner: SendAll<Sender<Result<S::Item, S::Error>>, ResultStream<S>>
-}
-
-struct ResultStream<S: Stream>(S);
-
-/// Spawns a `stream` onto the instance of `Executor` provided, `executor`,
-/// returning a handle representing the remote stream.
-///
-/// The `stream` will be canceled if the `SpawnHandle` is dropped.
-///
-/// The `SpawnHandle` returned is a stream that is a proxy for `stream` itself.
-/// When `stream` has additional items available, then the `SpawnHandle`
-/// will have those same items available.
-///
-/// At most `buffer + 1` elements will be buffered at a time. If the buffer
-/// is full, then `stream` will stop progressing until more space is available.
-/// This allows the `SpawnHandle` to exert backpressure on the `stream`.
-///
-/// # Panics
-///
-/// This function will panic if `executor` is unable spawn a `Future` containing
-/// the entirety of the `stream`.
-pub fn spawn<S, E>(stream: S, executor: &E, buffer: usize) -> SpawnHandle<S::Item, S::Error>
-    where S: Stream,
-          E: Executor<Execute<S>>
-{
-    let (tx, rx) = channel(buffer);
-    executor.execute(Execute {
-        inner: tx.send_all(ResultStream(stream))
-    }).expect("failed to spawn stream");
-    SpawnHandle {
-        rx: rx
-    }
-}
-
-/// Spawns a `stream` onto the instance of `Executor` provided, `executor`,
-/// returning a handle representing the remote stream, with unbounded buffering.
-///
-/// The `stream` will be canceled if the `SpawnHandle` is dropped.
-///
-/// The `SpawnHandle` returned is a stream that is a proxy for `stream` itself.
-/// When `stream` has additional items available, then the `SpawnHandle`
-/// will have those same items available.
-///
-/// An unbounded buffer is used, which means that values will be buffered as
-/// fast as `stream` can produce them, without any backpressure. Therefore, if
-/// `stream` is an infinite stream, it can use an unbounded amount of memory, and
-/// potentially hog CPU resources.
-///
-/// # Panics
-///
-/// This function will panic if `executor` is unable spawn a `Future` containing
-/// the entirety of the `stream`.
-pub fn spawn_unbounded<S, E>(stream: S, executor: &E) -> SpawnHandle<S::Item, S::Error>
-    where S: Stream,
-          E: Executor<Execute<S>>
-{
-    let (tx, rx) = channel2(None);
-    executor.execute(Execute {
-        inner: tx.send_all(ResultStream(stream))
-    }).expect("failed to spawn stream");
-    SpawnHandle {
-        rx: rx
-    }
-}
-
-impl<I, E> Stream for SpawnHandle<I, E> {
-    type Item = I;
-    type Error = E;
-
-    fn poll(&mut self) -> Poll<Option<I>, E> {
-        match self.rx.poll() {
-            Ok(Async::Ready(Some(Ok(t)))) => Ok(Async::Ready(Some(t.into()))),
-            Ok(Async::Ready(Some(Err(e)))) => Err(e),
-            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
-            Ok(Async::NotReady) => Ok(Async::NotReady),
-            Err(_) => unreachable!("mpsc::Receiver should never return Err"),
-        }
-    }
-}
-
-impl<I, E> fmt::Debug for SpawnHandle<I, E> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.debug_struct("SpawnHandle")
-         .finish()
-    }
-}
-
-impl<S: Stream> Future for Execute<S> {
-    type Item = ();
-    type Error = ();
-
-    fn poll(&mut self) -> Poll<(), ()> {
-        match self.inner.poll() {
-            Ok(Async::NotReady) => Ok(Async::NotReady),
-            _ => Ok(Async::Ready(()))
-        }
-    }
-}
-
-impl<S: Stream> fmt::Debug for Execute<S> {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.debug_struct("Execute")
-         .finish()
-    }
-}
-
-impl<S: Stream> Stream for ResultStream<S> {
-    type Item = Result<S::Item, S::Error>;
-    type Error = SendError<Self::Item>;
-    fn poll(&mut self) -> Poll<Option<Result<S::Item, S::Error>>, Self::Error> {
-        match self.0.poll() {
-            Ok(Async::Ready(Some(item))) => Ok(Async::Ready(Some(Ok(item)))),
-            Err(e) => Ok(Async::Ready(Some(Err(e)))),
-            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
-            Ok(Async::NotReady) => Ok(Async::NotReady)
-        }
-    }
-}
-
 /*
  *
  * ===== impl Inner =====


Test patch:
diff --git a/tests/mpsc.rs b/tests/mpsc.rs
index a4e114dc..a411bcef 100644
--- a/tests/mpsc.rs
+++ b/tests/mpsc.rs
@@ -3,8 +3,7 @@
 extern crate futures;
 
 use futures::{Future, Stream, Sink, Async, AsyncSink};
-use futures::future::{lazy, ok};
-use futures::stream::unfold;
+use futures::future::lazy;
 use futures::sync::mpsc;
 
 use std::time::Duration;
@@ -134,15 +133,6 @@ fn tx_close_gets_none() {
     }).wait().unwrap();
 }
 
-#[test]
-fn spawn_sends_items() {
-    let core = local_executor::Core::new();
-    let stream = unfold(0, |i| Some(ok::<_,u8>((i, i + 1))));
-    let rx = mpsc::spawn(stream, &core, 1);
-    assert_eq!(core.run(rx.take(4).collect()).unwrap(),
-               [0, 1, 2, 3]);
-}
-
 #[test]
 fn stress_shared_unbounded() {
     const AMT: u32 = 10000;
diff --git a/tests/stream.rs b/tests/stream.rs
index 0d5961a5..cd661300 100644
--- a/tests/stream.rs
+++ b/tests/stream.rs
@@ -1,10 +1,10 @@
 #[macro_use]
 extern crate futures;
 
-use futures::{Poll, Future, Stream, Sink};
+use futures::{Async, Future, Poll, Sink, Stream};
 use futures::executor;
-use futures::future::{ok, err};
-use futures::stream::{empty, iter, Peekable, BoxStream};
+use futures::future::{err, ok};
+use futures::stream::{empty, iter, poll_fn, BoxStream, Peekable};
 use futures::sync::oneshot;
 use futures::sync::mpsc;
 
@@ -356,3 +356,18 @@ fn concat2() {
     let c = empty::<Vec<()>, ()>();
     assert_done(move || c.concat2(), Ok(vec![]))
 }
+
+#[test]
+fn stream_poll_fn() {
+    let mut counter = 5usize;
+
+    let read_stream = poll_fn(move || -> Poll<Option<usize>, std::io::Error> {
+        if counter == 0 {
+            return Ok(Async::Ready(None));
+        }
+        counter -= 1;
+        Ok(Async::Ready(Some(counter)))
+    });
+
+    assert_eq!(read_stream.wait().count(), 5);
+}