stream: stop polling the underlying stream once map_while yields None (#8233)

tokio-stream/src/stream_ext/map_while.rs

@@ -3,6 +3,7 @@ use crate::Stream;
 use core::fmt;
 use core::pin::Pin;
 use core::task::{Context, Poll};
+use futures_core::FusedStream;
 use pin_project_lite::pin_project;
 
 pin_project! {
@@ -12,6 +13,7 @@ pin_project! {
         #[pin]
         stream: St,
         f: F,
+        done: bool,
     }
 }
 
@@ -22,13 +24,18 @@ where
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("MapWhile")
             .field("stream", &self.stream)
+            .field("done", &self.done)
             .finish()
     }
 }
 
 impl<St, F> MapWhile<St, F> {
     pub(super) fn new(stream: St, f: F) -> Self {
-        MapWhile { stream, f }
+        MapWhile {
+            stream,
+            f,
+            done: false,
+        }
     }
 }
 
@@ -41,12 +48,36 @@ where
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
         let me = self.project();
+        if *me.done {
+            return Poll::Ready(None);
+        }
+
         let f = me.f;
-        me.stream.poll_next(cx).map(|opt| opt.and_then(f))
+        let done = me.done;
+        me.stream.poll_next(cx).map(|opt| {
+            let mapped = opt.and_then(f);
+            if mapped.is_none() {
+                *done = true;
+            }
+            mapped
+        })
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.done {
+            return (0, Some(0));
+        }
+
         let (_, upper) = self.stream.size_hint();
         (0, upper)
     }
 }
+
+impl<St, F> FusedStream for MapWhile<St, F>
+where
+    Self: Stream,
+{
+    fn is_terminated(&self) -> bool {
+        self.done
+    }
+}

tokio-stream/tests/stream_fused.rs

@@ -129,6 +129,26 @@ async fn take_while_terminated_after_predicate_fails() {
     assert!(stream.is_terminated());
 }
 
+// ── map_while ─────────────────────────────────────────────────────────────────
+
+#[tokio::test]
+async fn map_while_not_terminated_before_closure_returns_none() {
+    let stream =
+        tokio_stream::iter(vec![1, 2, 3]).map_while(|x| if x < 10 { Some(x) } else { None });
+    assert!(!stream.is_terminated());
+}
+
+#[tokio::test]
+async fn map_while_terminated_after_closure_returns_none() {
+    let mut stream =
+        tokio_stream::iter(vec![1, 5, 2]).map_while(|x| if x < 3 { Some(x) } else { None });
+    assert_eq!(stream.next().await, Some(1));
+    assert!(!stream.is_terminated());
+    // closure returns `None` on 5 → done flag set
+    assert_eq!(stream.next().await, None);
+    assert!(stream.is_terminated());
+}
+
 // ── then ─────────────────────────────────────────────────────────────────────
 
 #[tokio::test]

tokio-stream/tests/stream_map_while.rs

@@ -0,0 +1,22 @@
+use tokio_stream::StreamExt;
+
+#[tokio::test]
+async fn map_while_yields_until_closure_returns_none() {
+    let mut stream =
+        tokio_stream::iter(1..=10).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);
+}
+
+#[tokio::test]
+async fn map_while_does_not_poll_after_closure_returns_none() {
+    // Once the closure returns `None`, the underlying stream must not be polled
+    // again, so the trailing `2` is never yielded.
+    let mut stream =
+        tokio_stream::iter(vec![1, 5, 2]).map_while(|x| if x < 3 { Some(x) } else { None });
+    assert_eq!(stream.next().await, Some(1));
+    assert_eq!(stream.next().await, None);
+    assert_eq!(stream.next().await, None);
+}