adds stream::find_map combinator

montekki · montekki · commit 45bd0ef13df4 · 2019-09-10T09:59:00.000+03:00
diff --git a/src/stream/stream/find_map.rs b/src/stream/stream/find_map.rs
@@ -0,0 +1,50 @@
+use std::marker::PhantomData;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+#[allow(missing_debug_implementations)]
+pub struct FindMapFuture<'a, S, F, T, B> {
+    stream: &'a mut S,
+    f: F,
+    __b: PhantomData<B>,
+    __t: PhantomData<T>,
+}
+
+impl<'a, S, B, F, T> FindMapFuture<'a, S, F, T, B> {
+    pin_utils::unsafe_pinned!(stream: &'a mut S);
+    pin_utils::unsafe_unpinned!(f: F);
+
+    pub(super) fn new(stream: &'a mut S, f: F) -> Self {
+        FindMapFuture {
+            stream,
+            f,
+            __b: PhantomData,
+            __t: PhantomData,
+        }
+    }
+}
+
+impl<'a, S, B, F> futures_core::future::Future for FindMapFuture<'a, S, F, S::Item, B>
+where
+    S: futures_core::stream::Stream + Unpin + Sized,
+    F: FnMut(S::Item) -> Option<B>,
+{
+    type Output = Option<B>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        use futures_core::stream::Stream;
+
+        let item = futures_core::ready!(self.as_mut().stream().poll_next(cx));
+
+        match item {
+            Some(v) => match (self.as_mut().f())(v) {
+                Some(v) => Poll::Ready(Some(v)),
+                None => {
+                    cx.waker().wake_by_ref();
+                    Poll::Pending
+                }
+            },
+            None => Poll::Ready(None),
+        }
+    }
+}
diff --git a/src/stream/stream/mod.rs b/src/stream/stream/mod.rs
@@ -23,6 +23,7 @@
 
 mod all;
 mod any;
+mod find_map;
 mod min_by;
 mod next;
 mod take;
@@ -31,6 +32,7 @@ pub use take::Take;
 
 use all::AllFuture;
 use any::AnyFuture;
+use find_map::FindMapFuture;
 use min_by::MinByFuture;
 use next::NextFuture;
 
@@ -48,13 +50,15 @@ cfg_if! {
             ($a:lifetime, $f:tt, $o:ty) => (ImplFuture<$a, $o>);
             ($a:lifetime, $f:tt, $o:ty, $t1:ty) => (ImplFuture<$a, $o>);
             ($a:lifetime, $f:tt, $o:ty, $t1:ty, $t2:ty) => (ImplFuture<$a, $o>);
+            ($a:lifetime, $f:tt, $o:ty, $t1:ty, $t2:ty, $t3:ty) => (ImplFuture<$a, $o>);
 
         }
     } else {
         macro_rules! ret {
             ($a:lifetime, $f:tt, $o:ty) => ($f<$a, Self>);
             ($a:lifetime, $f:tt, $o:ty, $t1:ty) => ($f<$a, Self, $t1>);
             ($a:lifetime, $f:tt, $o:ty, $t1:ty, $t2:ty) => ($f<$a, Self, $t1, $t2>);
+            ($a:lifetime, $f:tt, $o:ty, $t1:ty, $t2:ty, $t3:ty) => ($f<$a, Self, $t1, $t2, $t3>);
         }
     }
 }
@@ -218,6 +222,29 @@ pub trait Stream {
         }
     }
 
+    /// Applies function to the elements of stream and returns the first non-none result.
+    ///
+    /// ```
+    /// # fn main() { async_std::task::block_on(async {
+    /// #
+    /// use async_std::prelude::*;
+    /// use std::collections::VecDeque;
+    ///
+    /// let mut s: VecDeque<&str> = vec!["lol", "NaN", "2", "5"].into_iter().collect();
+    /// let first_number = s.find_map(|s| s.parse().ok()).await;
+    ///
+    /// assert_eq!(first_number, Some(2));
+    /// #
+    /// # }) }
+    /// ```
+    fn find_map<F, B>(&mut self, f: F) -> ret!('_, FindMapFuture, Option<B>, F, Self::Item, B)
+    where
+        Self: Sized,
+        F: FnMut(Self::Item) -> Option<B>,
+    {
+        FindMapFuture::new(self, f)
+    }
+
     /// Tests if any element of the stream matches a predicate.
     ///
     /// `any()` takes a closure that returns `true` or `false`. It applies
