feat: Add StreamExt::switch

futures-util/src/stream/mod.rs

@@ -20,8 +20,8 @@ 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,
+    SelectNextSome, Skip, SkipWhile, StreamExt, StreamFuture, Switch, Take, TakeUntil, TakeWhile,
+    Then, TryFold, TryForEach, Unzip, Zip,
 };
 
 #[cfg(feature = "std")]

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

@@ -122,6 +122,16 @@ pub use self::skip::Skip;
 mod skip_while;
 pub use self::skip_while::SkipWhile;
 
+mod switch;
+
+delegate_all!(
+    /// Stream for the [`switch`](StreamExt::switch) method.
+    Switch<St>(
+        switch::Switch<St, St::Item>
+    ): Debug + Sink + Stream + FusedStream + AccessInner[St, (.)] + New[|x: St| switch::Switch::new(x)]
+    where St: Stream
+);
+
 mod take;
 pub use self::take::Take;
 
@@ -1140,6 +1150,108 @@ pub trait StreamExt: Stream {
         assert_future::<(), _>(ForEachConcurrent::new(self, limit.into(), f))
     }
 
+    /// Flattens a higher-order stream into a first-order stream.
+    ///
+    /// This combinator flattens a stream of streams, i.e. an outer stream
+    /// yielding inner streams. This combinator always keeps the most recently
+    /// yielded inner stream, and yields items from it, until the outer stream
+    /// produces a new inner stream, at which point the inner stream to yield
+    /// items from is switched to the new one.
+    ///
+    /// # Examples
+    ///
+    /// An empty (outer) stream can be switched:
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::stream::{self, Empty, StreamExt, Switch};
+    ///
+    /// let stream: Switch<Empty<Empty<()>>> = stream::empty().switch();
+    ///
+    /// assert!(stream.collect::<Vec<_>>().await.is_empty());
+    /// # });
+    /// ```
+    ///
+    /// An outer stream can produce several inner streams — once switched, the
+    /// last one is immediately selected:
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::stream::{self, StreamExt};
+    ///
+    /// let stream = stream::iter([
+    ///     stream::iter([1, 2, 3]),
+    ///     stream::iter([4, 5, 6]),
+    ///     stream::iter([7, 8, 9]),
+    /// ])
+    /// .switch();
+    ///
+    /// assert_eq!(vec![7, 8, 9], stream.collect::<Vec<_>>().await);
+    /// # });
+    /// ```
+    ///
+    /// One of the most interesting usecase is when an outer stream produces new
+    /// inner  streams dynamically. Let's imagine the outer stream produces inner
+    /// streams yielding a sequence of integers starting from a particular value.
+    /// For example:
+    ///
+    /// - outer stream yields 7:
+    ///   - inner stream yields 7, 8, 9, 10, 11…
+    /// - outer stream yields 42:
+    ///   - inner stream yields 42, 43, 44, 45, 46…
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::channel::mpsc;
+    /// use futures::stream::{self, StreamExt};
+    /// use futures::task::Poll;
+    /// use std::pin::pin;
+    ///
+    /// // `receiver` is the outer stream: it yields integers received by
+    /// // `sender`.
+    /// let (sender, receiver) = mpsc::unbounded::<u8>();
+    ///
+    /// let mut stream = pin!(receiver
+    ///     // For every new received value from `sender`…
+    ///     .map(|init_value| {
+    ///         // … let's create a new inner stream.
+    ///         //
+    ///         // First off, the inner stream starts with `init_value`. Then,
+    ///         // it continues with incremented integers.
+    ///         let mut next_value = init_value;
+    ///
+    ///         stream::poll_fn(move |_| {
+    ///             let current_value = next_value;
+    ///             next_value += 1;
+    ///
+    ///             Poll::Ready(Some(current_value))
+    ///         })
+    ///     })
+    ///     .switch());
+    ///
+    /// // `stream` is pending until the outer stream yields something.
+    /// sender.unbounded_send(7).unwrap();
+    ///
+    /// // `stream` has switched to the inner stream.
+    /// assert_eq!(vec![7, 8, 9, 10, 11], stream.by_ref().take(5).collect::<Vec<_>>().await);
+    /// assert_eq!(vec![12, 13, 14, 15, 16], stream.by_ref().take(5).collect::<Vec<_>>().await);
+    ///
+    /// // The outer stream will yield a new value, which will create a new
+    /// // inner stream.
+    /// sender.unbounded_send(42).unwrap();
+    ///
+    /// // `stream` has been “reset” and will produce a new inner stream.
+    /// assert_eq!(vec![42, 43, 44, 45, 46], stream.take(5).collect::<Vec<_>>().await);
+    /// # });
+    /// ```
+    fn switch(self) -> Switch<Self>
+    where
+        Self: Sized,
+        Self::Item: Stream,
+    {
+        assert_stream::<<Self::Item as Stream>::Item, _>(Switch::new(self))
+    }
+
     /// Attempt to execute an accumulating asynchronous computation over a
     /// stream, collecting all the values into one final result.
     ///

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

@@ -0,0 +1,90 @@
+use core::pin::Pin;
+
+use futures_core::stream::Stream;
+use futures_core::task::{Context, Poll};
+use futures_core::FusedStream;
+use pin_project_lite::pin_project;
+
+pin_project! {
+    /// Stream for the [`switch`](super::StreamExt::switch) method.
+    #[derive(Debug)]
+    #[must_use = "futures do nothing unless you `.await` or poll them"]
+    pub struct Switch<Outer, Inner> {
+        #[pin]
+        outer_stream: Outer,
+        outer_stream_is_closed: bool,
+        #[pin]
+        inner_stream: Option<Inner>,
+    }
+}
+
+impl<Outer: Stream> Switch<Outer, Outer::Item> {
+    /// Creates a new [`Switch`] stream.
+    pub(super) fn new(outer_stream: Outer) -> Self {
+        Self { outer_stream, outer_stream_is_closed: false, inner_stream: None }
+    }
+
+    delegate_access_inner!(outer_stream, Outer, ());
+}
+
+impl<Outer> FusedStream for Switch<Outer, Outer::Item>
+where
+    Outer: FusedStream,
+    Outer::Item: Stream,
+{
+    fn is_terminated(&self) -> bool {
+        self.inner_stream.is_none() && self.outer_stream.is_terminated()
+    }
+}
+
+impl<Outer> Stream for Switch<Outer, Outer::Item>
+where
+    Outer: Stream,
+    Outer::Item: Stream,
+{
+    type Item = <Outer::Item as Stream>::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+
+        // Poll the latest inner stream eagerly.
+        if !*this.outer_stream_is_closed {
+            while let Poll::Ready(ready) = this.outer_stream.as_mut().poll_next(cx) {
+                match ready {
+                    Some(inner_stream) => {
+                        this.inner_stream.set(Some(inner_stream));
+                    }
+
+                    None => {
+                        *this.outer_stream_is_closed = true;
+                        break;
+                    }
+                }
+            }
+        }
+
+        match this.inner_stream.as_mut().as_pin_mut() {
+            // No inner stream has been produced yet.
+            None => {
+                // The stream' state is the outer stream' state.
+                if *this.outer_stream_is_closed {
+                    Poll::Ready(None)
+                } else {
+                    Poll::Pending
+                }
+            }
+
+            // An inner stream exists: poll it!
+            Some(inner_stream) => match inner_stream.poll_next(cx) {
+                // Inner stream produced an item.
+                Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
+
+                // Both inner and outer streams are closed.
+                Poll::Ready(None) if *this.outer_stream_is_closed => Poll::Ready(None),
+
+                // Only inner stream is closed or is pending.
+                Poll::Ready(None) | Poll::Pending => Poll::Pending,
+            },
+        }
+    }
+}

futures/tests/stream.rs

@@ -5,6 +5,7 @@ use std::rc::Rc;
 use std::sync::Arc;
 use std::task::Context;
 
+use assert_matches::assert_matches;
 use futures::channel::mpsc;
 use futures::executor::block_on;
 use futures::future::{self, Future};
@@ -592,3 +593,65 @@ fn any() {
         assert!(!any);
     });
 }
+
+#[test]
+fn switch() {
+    use std::pin::pin;
+
+    block_on(async {
+        // Empty.
+        let stream = stream::empty::<stream::Empty<stream::Empty<()>>>().switch();
+        assert!(stream.collect::<Vec<_>>().await.is_empty());
+
+        // Stream is the last inner stream.
+        let stream = stream::iter([stream::iter([1, 2, 3]), stream::iter([4, 5, 6])]).switch();
+        assert_eq!(stream.collect::<Vec<_>>().await, vec![4, 5, 6]);
+
+        // Outer stream is closed: stream is closed.
+        let mut stream = stream::poll_fn(|_| Poll::Ready(None::<stream::Empty<()>>)).switch();
+        assert_matches!(stream.next().now_or_never(), Some(None));
+
+        // Outer stream is pending: stream is pending.
+        let mut stream = stream::poll_fn(|_| Poll::<Option<stream::Empty<()>>>::Pending).switch();
+        assert_matches!(stream.next().now_or_never(), None);
+
+        // Outer stream is closed after first poll, and inner stream is closed: stream is closed.
+        let mut stream =
+            pin!(stream::once(async { stream::poll_fn(|_| Poll::Ready(None::<()>)) }).switch());
+        assert_matches!(stream.next().now_or_never(), Some(None));
+
+        // Inner stream is closed: stream is pending.
+        let mut stream = {
+            let mut yielded_once = false;
+
+            stream::poll_fn(move |_| {
+                if yielded_once {
+                    Poll::Pending
+                } else {
+                    yielded_once = true;
+
+                    Poll::Ready(Some(stream::poll_fn(|_| Poll::Ready(None::<()>))))
+                }
+            })
+            .switch()
+        };
+        assert_matches!(stream.next().now_or_never(), None);
+
+        // Inner stream is pending: stream is pending.
+        let mut stream = {
+            let mut yielded_once = false;
+
+            stream::poll_fn(move |_| {
+                if yielded_once {
+                    Poll::Pending
+                } else {
+                    yielded_once = true;
+
+                    Poll::Ready(Some(stream::pending::<()>()))
+                }
+            })
+            .switch()
+        };
+        assert_matches!(stream.next().now_or_never(), None);
+    });
+}