Add StreamExt::enumerate

Author: taiki-e

futures-util/src/stream/enumerate.rs

@@ -0,0 +1,61 @@
+use core::pin::Pin;
+use futures_core::stream::{FusedStream, Stream};
+use futures_core::task::{Context, Poll};
+use futures_sink::Sink;
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+
+/// Stream for the [`enumerate`](super::EnumerateExt::enumerate) method.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct Enumerate<St: Stream> {
+    stream: St,
+    count: usize,
+}
+
+impl<St: Stream + Unpin> Unpin for Enumerate<St> {}
+
+impl<St: Stream> Enumerate<St> {
+    unsafe_pinned!(stream: St);
+    unsafe_unpinned!(count: usize);
+
+    pub(super) fn new(stream: St) -> Enumerate<St> {
+        Enumerate {
+            stream,
+            count: 0,
+        }
+    }
+}
+
+impl<St: Stream + FusedStream> FusedStream for Enumerate<St> {
+    fn is_terminated(&self) -> bool {
+        self.stream.is_terminated()
+    }
+}
+
+impl<St: Stream> Stream for Enumerate<St> {
+    type Item = (usize, St::Item);
+
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Self::Item>> {
+        match ready!(self.as_mut().stream().poll_next(cx)) {
+            Some(item) => {
+                let count = self.count;
+                *self.as_mut().count() += 1;
+                Poll::Ready(Some((count, item)))
+            }
+            None => Poll::Ready(None),
+        }
+    }
+}
+
+// Forwarding impl of Sink from the underlying stream
+impl<S, Item> Sink<Item> for Enumerate<S>
+where
+    S: Stream + Sink<Item>,
+{
+    type SinkError = S::SinkError;
+
+    delegate_sink!(stream, Item);
+}

futures-util/src/stream/mod.rs

@@ -30,6 +30,9 @@ pub use self::concat::Concat;
 mod empty;
 pub use self::empty::{empty, Empty};
 
+mod enumerate;
+pub use self::enumerate::Enumerate;
+
 mod filter;
 pub use self::filter::Filter;
 
@@ -242,6 +245,51 @@ pub trait StreamExt: Stream {
         Map::new(self, f)
     }
 
+    /// Creates a stream which gives the current iteration count as well as
+    /// the next value.
+    ///
+    /// The stream returned yields pairs `(i, val)`, where `i` is the
+    /// current index of iteration and `val` is the value returned by the
+    /// stream.
+    ///
+    /// `enumerate()` keeps its count as a [`usize`]. If you want to count by a
+    /// different sized integer, the [`zip`](StreamExt::zip) function provides similar
+    /// functionality.
+    ///
+    /// # Overflow Behavior
+    ///
+    /// The method does no guarding against overflows, so enumerating more than
+    /// [`usize::max_value()`] elements either produces the wrong result or panics. If
+    /// debug assertions are enabled, a panic is guaranteed.
+    ///
+    /// # Panics
+    ///
+    /// The returned stream might panic if the to-be-returned index would
+    /// overflow a [`usize`].
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await, await_macro, futures_api)]
+    /// # futures::executor::block_on(async {
+    /// use futures::stream::{self, StreamExt};
+    ///
+    /// let stream = stream::iter(vec!['a', 'b', 'c']);
+    ///
+    /// let mut stream = stream.enumerate();
+    ///
+    /// assert_eq!(await!(stream.next()), Some((0, 'a')));
+    /// assert_eq!(await!(stream.next()), Some((1, 'b')));
+    /// assert_eq!(await!(stream.next()), Some((2, 'c')));
+    /// assert_eq!(await!(stream.next()), None);
+    /// # });
+    /// ```
+    fn enumerate(self) -> Enumerate<Self>
+        where Self: Sized,
+    {
+        Enumerate::new(self)
+    }
+
     /// Filters the values produced by this stream according to the provided
     /// asynchronous predicate.
     ///

futures/src/lib.rs

@@ -345,9 +345,10 @@ pub mod stream {
         unfold, Unfold,
 
         StreamExt,
-        Chain, Collect, Concat, Filter, FilterMap, Flatten, Fold, Forward,
-        ForEach, Fuse, StreamFuture, Inspect, Map, Next, SelectNextSome,
-        Peekable, Select, Skip, SkipWhile, Take, TakeWhile, Then, Zip
+        Chain, Collect, Concat, Enumerate, Filter, FilterMap, Flatten, Fold,
+        Forward, ForEach, Fuse, StreamFuture, Inspect, Map, Next,
+        SelectNextSome, Peekable, Select, Skip, SkipWhile, Take, TakeWhile,
+        Then, Zip
     };
 
     #[cfg(feature = "alloc")]