Add StreamExt::cycle

Author: taiki-e

futures-util/src/stream/iter.rs

@@ -3,7 +3,7 @@ use futures_core::stream::Stream;
 use futures_core::task::{Context, Poll};
 
 /// Stream for the [`iter`] function.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 #[must_use = "streams do nothing unless polled"]
 pub struct Iter<I> {
     iter: I,

futures-util/src/stream/mod.rs

@@ -13,8 +13,8 @@ pub use futures_core::stream::{FusedStream, Stream, TryStream};
 #[allow(clippy::module_inception)]
 mod stream;
 pub use self::stream::{
-    Chain, Collect, Concat, Enumerate, Filter, FilterMap, FlatMap, Flatten, Fold, ForEach, Fuse,
-    Inspect, Map, Next, Peek, Peekable, Scan, SelectNextSome, Skip, SkipWhile, StreamExt,
+    Chain, Collect, Concat, Cycle, Enumerate, Filter, FilterMap, FlatMap, Flatten, Fold, ForEach,
+    Fuse, Inspect, Map, Next, Peek, Peekable, Scan, SelectNextSome, Skip, SkipWhile, StreamExt,
     StreamFuture, Take, TakeWhile, TakeUntil, Then, Zip,
 };
 

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

@@ -0,0 +1,69 @@
+use core::pin::Pin;
+use core::usize;
+use futures_core::stream::{FusedStream, Stream};
+use futures_core::task::{Context, Poll};
+use pin_project::pin_project;
+
+/// Stream for the [`cycle`](super::StreamExt::cycle) method.
+#[pin_project]
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct Cycle<St> {
+    orig: St,
+    #[pin]
+    stream: St,
+}
+
+impl<St> Cycle<St>
+where
+    St: Clone + Stream,
+{
+    pub(super) fn new(stream: St) -> Self {
+        Self {
+            orig: stream.clone(),
+            stream,
+        }
+    }
+}
+
+impl<St> Stream for Cycle<St>
+where
+    St: Clone + Stream,
+{
+    type Item = St::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+
+        match ready!(this.stream.as_mut().poll_next(cx)) {
+            None => {
+                this.stream.set(this.orig.clone());
+                this.stream.poll_next(cx)
+            }
+            item => Poll::Ready(item),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // the cycle stream is either empty or infinite
+        match self.orig.size_hint() {
+            size @ (0, Some(0)) => size,
+            (0, _) => (0, None),
+            _ => (usize::max_value(), None),
+        }
+    }
+}
+
+impl<St> FusedStream for Cycle<St>
+where
+    St: Clone + Stream,
+{
+    fn is_terminated(&self) -> bool {
+        // the cycle stream is either empty or infinite
+        if let (0, Some(0)) = self.size_hint() {
+            true
+        } else {
+            false
+        }
+    }
+}

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

@@ -33,6 +33,10 @@ mod concat;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::concat::Concat;
 
+mod cycle;
+#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
+pub use self::cycle::Cycle;
+
 mod enumerate;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::enumerate::Enumerate;
@@ -513,6 +517,36 @@ pub trait StreamExt: Stream {
         assert_future::<Self::Item, _>(Concat::new(self))
     }
 
+    /// Repeats a stream endlessly.
+    ///
+    /// The stream never terminates. Note that you likely want to avoid
+    /// usage of `collect` or such on the returned stream as it will exhaust
+    /// available memory as it tries to just fill up all RAM.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::stream::{self, StreamExt};
+    /// let a = [1, 2, 3];
+    /// let mut s = stream::iter(a.iter()).cycle();
+    ///
+    /// assert_eq!(s.next().await, Some(&1));
+    /// assert_eq!(s.next().await, Some(&2));
+    /// assert_eq!(s.next().await, Some(&3));
+    /// assert_eq!(s.next().await, Some(&1));
+    /// assert_eq!(s.next().await, Some(&2));
+    /// assert_eq!(s.next().await, Some(&3));
+    /// assert_eq!(s.next().await, Some(&1));
+    /// # });
+    /// ```
+    fn cycle(self) -> Cycle<Self>
+    where
+        Self: Sized + Clone,
+    {
+        assert_stream::<Self::Item, _>(Cycle::new(self))
+    }
+
     /// Execute an accumulating asynchronous computation over a stream,
     /// collecting all the values into one final result.
     ///

futures/src/lib.rs

@@ -456,7 +456,7 @@ pub mod stream {
         try_unfold, TryUnfold,
 
         StreamExt,
-        Chain, Collect, Concat, Enumerate, Filter, FilterMap, FlatMap, Flatten,
+        Chain, Collect, Concat, Cycle, Enumerate, Filter, FilterMap, FlatMap, Flatten,
         Fold, Forward, ForEach, Fuse, StreamFuture, Inspect, Map, Next,
         SelectNextSome, Peek, Peekable, Scan, Skip, SkipWhile, Take, TakeUntil,
         TakeWhile, Then, Zip,