From fb35c839443acd228ddefdd2a6985c07f98da86b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Braulio=20Valdivielso=20Mart=C3=ADnez?=
 <brlvldvlsmrtnz@gmail.com>
Date: Fri, 31 Dec 2021 13:53:09 +0000
Subject: [PATCH] tokio-stream: add `StreamExt::map_while` (#4351)

Fixes #4337

Rust 1.57 stabilized the `Iterator::map_while` API. This PR adds the
same functionality to the `StreamExt` trait, to keep parity.
---
 tokio-stream/src/stream_ext.rs           | 45 ++++++++++++++++++++
 tokio-stream/src/stream_ext/map_while.rs | 52 ++++++++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 tokio-stream/src/stream_ext/map_while.rs

diff --git a/tokio-stream/src/stream_ext.rs b/tokio-stream/src/stream_ext.rs
index cc4841e62..b79883bd6 100644
--- a/tokio-stream/src/stream_ext.rs
+++ b/tokio-stream/src/stream_ext.rs
@@ -28,6 +28,9 @@ use fuse::Fuse;
 mod map;
 use map::Map;
 
+mod map_while;
+use map_while::MapWhile;
+
 mod merge;
 use merge::Merge;
 
@@ -201,6 +204,48 @@ pub trait StreamExt: Stream {
         Map::new(self, f)
     }
 
+    /// Map this stream's items to a different type for as long as determined by
+    /// the provided closure. A stream of the target type will be returned,
+    /// which will yield elements until the closure returns `None`.
+    ///
+    /// The provided closure is executed over all elements of this stream as
+    /// they are made available, until it returns `None`. It is executed inline
+    /// with calls to [`poll_next`](Stream::poll_next). Once `None` is returned,
+    /// the underlying stream will not be polled again.
+    ///
+    /// Note that this function consumes the stream passed into it and returns a
+    /// wrapped version of it, similar to the [`Iterator::map_while`] method in the
+    /// standard library.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// use tokio_stream::{self as stream, StreamExt};
+    ///
+    /// let stream = stream::iter(1..=10);
+    /// let mut stream = stream.map_while(|x| {
+    ///     if x < 4 {
+    ///         Some(x + 3)
+    ///     } else {
+    ///         None
+    ///     }
+    /// });
+    /// assert_eq!(stream.next().await, Some(4));
+    /// assert_eq!(stream.next().await, Some(5));
+    /// assert_eq!(stream.next().await, Some(6));
+    /// assert_eq!(stream.next().await, None);
+    /// # }
+    /// ```
+    fn map_while<T, F>(self, f: F) -> MapWhile<Self, F>
+    where
+        F: FnMut(Self::Item) -> Option<T>,
+        Self: Sized,
+    {
+        MapWhile::new(self, f)
+    }
+
     /// Maps this stream's items asynchronously to a different type, returning a
     /// new stream of the resulting type.
     ///
diff --git a/tokio-stream/src/stream_ext/map_while.rs b/tokio-stream/src/stream_ext/map_while.rs
new file mode 100644
index 000000000..d4fd82565
--- /dev/null
+++ b/tokio-stream/src/stream_ext/map_while.rs
@@ -0,0 +1,52 @@
+use crate::Stream;
+
+use core::fmt;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+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, F> {
+        #[pin]
+        stream: St,
+        f: F,
+    }
+}
+
+impl<St, F> fmt::Debug for MapWhile<St, F>
+where
+    St: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("MapWhile")
+            .field("stream", &self.stream)
+            .finish()
+    }
+}
+
+impl<St, F> MapWhile<St, F> {
+    pub(super) fn new(stream: St, f: F) -> Self {
+        MapWhile { stream, f }
+    }
+}
+
+impl<St, F, T> Stream for MapWhile<St, F>
+where
+    St: Stream,
+    F: FnMut(St::Item) -> Option<T>,
+{
+    type Item = T;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
+        let me = self.project();
+        let f = me.f;
+        me.stream.poll_next(cx).map(|opt| opt.and_then(f))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let (_, upper) = self.stream.size_hint();
+        (0, upper)
+    }
+}