Impl MapWhile

futures-util/src/stream/mod.rs

@@ -19,9 +19,9 @@ pub use futures_core::stream::{FusedStream, Stream, TryStream};
 mod stream;
 pub use self::stream::{
     All, Any, Chain, Collect, Concat, Count, Cycle, Enumerate, Filter, FilterMap, FlatMap, Flatten,
-    Fold, ForEach, Fuse, Inspect, Map, Next, NextIf, NextIfEq, Peek, PeekMut, Peekable, Scan,
-    SelectNextSome, Skip, SkipWhile, StreamExt, StreamFuture, Take, TakeUntil, TakeWhile, Then,
-    TryFold, TryForEach, Unzip, Zip,
+    Fold, ForEach, Fuse, Inspect, Map, MapWhile, Next, NextIf, NextIfEq, Peek, PeekMut, Peekable,
+    Scan, SelectNextSome, Skip, SkipWhile, StreamExt, StreamFuture, Take, TakeUntil, TakeWhile,
+    Then, TryFold, TryForEach, Unzip, Zip,
 };
 
 #[cfg(feature = "std")]

futures-util/src/stream/stream/map_while.rs

@@ -0,0 +1,115 @@
+use core::fmt;
+use core::pin::Pin;
+use futures_core::future::Future;
+use futures_core::ready;
+use futures_core::stream::{FusedStream, Stream};
+use futures_core::task::{Context, Poll};
+#[cfg(feature = "sink")]
+use futures_sink::Sink;
+use pin_project_lite::pin_project;
+
+pin_project! {
+    /// Stream for the [`map_while`](super::StreamExt::map_while) method.
+    #[must_use = "streams do nothing unless polled"]
+    pub struct MapWhile<St: Stream, Fut, F> {
+        #[pin]
+        stream: St,
+        f: F,
+        #[pin]
+        pending_fut: Option<Fut>,
+        done_mapping: bool,
+    }
+}
+
+impl<St, Fut, F> fmt::Debug for MapWhile<St, Fut, F>
+where
+    St: Stream + fmt::Debug,
+    St::Item: fmt::Debug,
+    Fut: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("MapWhile")
+            .field("stream", &self.stream)
+            .field("pending_fut", &self.pending_fut)
+            .field("done_mapping", &self.done_mapping)
+            .finish()
+    }
+}
+
+impl<St, Fut, F, T> MapWhile<St, Fut, F>
+where
+    St: Stream,
+    F: FnMut(St::Item) -> Fut,
+    Fut: Future<Output = Option<T>>,
+{
+    pub(super) fn new(stream: St, f: F) -> Self {
+        Self { stream, f, pending_fut: None, done_mapping: false }
+    }
+
+    delegate_access_inner!(stream, St, ());
+}
+
+impl<St, Fut, F, T> Stream for MapWhile<St, Fut, F>
+where
+    St: Stream,
+    F: FnMut(St::Item) -> Fut,
+    Fut: Future<Output = Option<T>>,
+{
+    type Item = T;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
+        if self.done_mapping {
+            return Poll::Ready(None);
+        }
+
+        let mut this = self.project();
+
+        Poll::Ready(loop {
+            if let Some(fut) = this.pending_fut.as_mut().as_pin_mut() {
+                let mapped = ready!(fut.poll(cx));
+                this.pending_fut.set(None);
+
+                if mapped.is_none() {
+                    *this.done_mapping = true;
+                }
+
+                break mapped;
+            } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) {
+                this.pending_fut.set(Some((this.f)(item)));
+            } else {
+                break None;
+            }
+        })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.done_mapping {
+            return (0, Some(0));
+        }
+
+        let (_, upper) = self.stream.size_hint();
+        (0, upper) // can't know a lower bound, due to the predicate
+    }
+}
+
+impl<St, Fut, F, T> FusedStream for MapWhile<St, Fut, F>
+where
+    St: FusedStream,
+    F: FnMut(St::Item) -> Fut,
+    Fut: Future<Output = Option<T>>,
+{
+    fn is_terminated(&self) -> bool {
+        self.done_mapping || self.stream.is_terminated()
+    }
+}
+
+// Forwarding impl of Sink from the underlying stream
+#[cfg(feature = "sink")]
+impl<S, Fut, F, Item> Sink<Item> for MapWhile<S, Fut, F>
+where
+    S: Stream + Sink<Item>,
+{
+    type Error = S::Error;
+
+    delegate_sink!(stream, Item);
+}

futures-util/src/stream/stream/mod.rs

@@ -128,6 +128,9 @@ pub use self::take::Take;
 mod take_while;
 pub use self::take_while::TakeWhile;
 
+mod map_while;
+pub use self::map_while::MapWhile;
+
 mod take_until;
 pub use self::take_until::TakeUntil;
 
@@ -994,6 +997,31 @@ pub trait StreamExt: Stream {
         assert_stream::<Self::Item, _>(TakeWhile::new(self, f))
     }
 
+    /// todo
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::future;
+    /// use futures::stream::{self, StreamExt};
+    ///
+    /// let stream = stream::iter(vec![1, 4, 0, 2]);
+    ///
+    /// let stream = stream.map_while(|x| future::ready(u32::checked_div(16, x)));
+    ///
+    /// assert_eq!(vec![16, 4], stream.collect::<Vec<_>>().await);
+    /// # });
+    /// ```
+    fn map_while<Fut, F, T>(self, f: F) -> MapWhile<Self, Fut, F>
+    where
+        F: FnMut(Self::Item) -> Fut,
+        Fut: Future<Output = Option<T>>,
+        Self: Sized,
+    {
+        assert_stream::<T, _>(MapWhile::new(self, f))
+    }
+
     /// Take elements from this stream until the provided future resolves.
     ///
     /// This function will take elements from the stream until the provided