add ParallelStream::find_map async-rs#2

Signed-off-by: Benjamin Coenen <[email protected]>

Author: bnjjj

src/par_stream/find_map.rs

@@ -0,0 +1,100 @@
+// use async_std::prelude::*;
+use async_std::future::Future;
+use async_std::sync::{self, Receiver};
+use async_std::task;
+
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use crate::ParallelStream;
+
+pin_project_lite::pin_project! {
+    /// A parallel stream that FindMaps value of another stream with a function.
+    #[derive(Debug)]
+    pub struct FindMap<T> {
+        #[pin]
+        receiver: Receiver<T>,
+        limit: Option<usize>,
+        already_sent: bool,
+    }
+}
+
+impl<T> FindMap<T>
+where
+    T: Send + 'static,
+{
+    /// Create a new instance of `FindMap`.
+    pub fn new<S, F, Fut>(mut stream: S, mut f: F) -> Self
+    where
+        S: ParallelStream,
+        F: FnMut(S::Item) -> Fut + Send + Sync + Copy + 'static,
+        Fut: Future<Output = Option<T>> + Send,
+    {
+        let (sender, receiver) = sync::channel(1);
+        let limit = stream.get_limit();
+        task::spawn(async move {
+            while let Some(item) = stream.next().await {
+                let sender = sender.clone();
+                task::spawn(async move {
+                    let res = f(item).await;
+                    if let Some(res) = res {
+                        sender.send(res).await;
+                    }
+                });
+            }
+        });
+        FindMap {
+            receiver,
+            limit,
+            already_sent: false,
+        }
+    }
+}
+
+impl<T> ParallelStream for FindMap<T>
+where
+    T: Send + 'static,
+{
+    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();
+        if *this.already_sent {
+            return Poll::Ready(None);
+        }
+        if let Poll::Ready(elt) = this.receiver.poll_next(cx) {
+            if let Some(elt) = elt {
+                *this.already_sent = true;
+                return Poll::Ready(Some(elt));
+            }
+            return Poll::Ready(None);
+        }
+        Poll::Pending
+    }
+
+    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 s = async_std::stream::from_iter(vec![1, 1, 2, 3, 2]);
+    let mut output: Vec<usize> = vec![];
+    let mut stream = crate::from_stream(s).find_map(|n| async move {
+        if n % 2 == 0 {
+            Some(42usize)
+        } else {
+            None
+        }
+    });
+    while let Some(n) = stream.next().await {
+        output.push(n);
+    }
+    assert_eq!(output, vec![42]);
+}

src/par_stream/mod.rs

@@ -5,11 +5,13 @@ use std::pin::Pin;
 
 use crate::FromParallelStream;
 
+pub use find_map::FindMap;
 pub use for_each::ForEach;
 pub use map::Map;
 pub use next::NextFuture;
 pub use take::Take;
 
+mod find_map;
 mod for_each;
 mod map;
 mod next;
@@ -29,6 +31,17 @@ pub trait ParallelStream: Sized + Send + Sync + Unpin + 'static {
     /// Get the max concurrency limit
     fn get_limit(&self) -> Option<usize>;
 
+    /// Applies `f` to each item of this stream in parallel, producing a new
+    /// stream with the results.
+    fn find_map<F, T, Fut>(self, f: F) -> FindMap<T>
+    where
+        F: FnMut(Self::Item) -> Fut + Send + Sync + Copy + 'static,
+        T: Send + 'static,
+        Fut: Future<Output = Option<T>> + Send,
+    {
+        FindMap::new(self, f)
+    }
+
     /// Applies `f` to each item of this stream in parallel, producing a new
     /// stream with the results.
     fn map<F, T, Fut>(self, f: F) -> Map<T>