Implementing from_fn

src/from_fn.rs

@@ -0,0 +1,91 @@
+use async_std::future::Future;
+use async_std::sync::{self, Receiver};
+use async_std::task;
+use async_std::task::{Context, Poll};
+use core::pin::Pin;
+use pin_project_lite::pin_project;
+
+use crate::ParallelStream;
+
+pin_project! {
+    /// A parallel stream that yields elements by calling a closure.
+    ///
+    /// This stream is created by the [`from_fn`] function.
+    /// See it documentation for more.
+    ///
+    /// [`from_fn`]: fn.from_fn.html
+    ///
+    /// # Examples
+    #[derive(Clone, Debug)]
+    pub struct FromFn<T, F> {
+        #[pin]
+        receiver: Receiver<T>,
+        f: F,
+        limit: Option<usize>,
+    }
+}
+
+/// Creates a parallel stream from a closure.
+pub fn from_fn<T, F, Fut>(mut f: F) -> FromFn<T, F>
+where
+    T: Send + Sync + Unpin + 'static,
+    F: FnMut() -> Fut + Send + Sync + Copy + 'static,
+    Fut: Future<Output = Option<T>> + Send,
+{
+    let (sender, receiver) = sync::channel(1);
+    task::spawn(async move {
+        let sender = sender.clone();
+        while let Some(val) = f().await {
+            sender.send(val).await;
+        }
+    });
+    FromFn {
+        f,
+        receiver,
+        limit: None,
+    }
+}
+
+impl<T: Send + Sync + Unpin + 'static, F, Fut> ParallelStream for FromFn<T, F>
+where
+    T: Send + Sync + Unpin + 'static,
+    F: FnMut() -> Fut + Send + Sync + Copy + 'static,
+    Fut: Future<Output = Option<T>> + Send,
+{
+    type Item = T;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        use async_std::prelude::*;
+        let this = self.project();
+        this.receiver.poll_next(cx)
+    }
+
+    fn limit(mut self, limit: impl Into<Option<usize>>) -> Self {
+        self.limit = limit.into();
+        self
+    }
+
+    fn get_limit(&self) -> Option<usize> {
+        self.limit
+    }
+}
+
+#[async_std::test]
+async fn smoke() {
+    let mut output = vec![];
+    let mut count = 0u8;
+    let mut stream = crate::from_fn(move || {
+        count += 1;
+        async move {
+            if count <= 3 {
+                Some(count)
+            } else {
+                None
+            }
+        }
+    });
+    while let Some(n) = stream.next().await {
+        output.push(n);
+    }
+    assert_eq!(output, vec![1, 2, 3]);
+}

src/lib.rs

@@ -44,11 +44,13 @@
 #![deny(missing_debug_implementations, nonstandard_style)]
 #![warn(missing_docs, missing_doc_code_examples)]
 
+mod from_fn;
 mod from_parallel_stream;
 mod from_stream;
 mod into_parallel_stream;
 mod par_stream;
 
+pub use from_fn::{from_fn, FromFn};
 pub use from_parallel_stream::FromParallelStream;
 pub use from_stream::{from_stream, FromStream};
 pub use into_parallel_stream::IntoParallelStream;