| 1 | use crate::Stream; |
| 2 | |
| 3 | use core::pin::Pin; |
| 4 | use core::task::{Context, Poll}; |
| 5 | |
| 6 | /// Stream for the [`iter`](fn@iter) function. |
| 7 | #[derive(Debug)] |
| 8 | #[must_use = "streams do nothing unless polled" ] |
| 9 | pub struct Iter<I> { |
| 10 | iter: I, |
| 11 | yield_amt: usize, |
| 12 | } |
| 13 | |
| 14 | impl<I> Unpin for Iter<I> {} |
| 15 | |
| 16 | /// Converts an `Iterator` into a `Stream` which is always ready |
| 17 | /// to yield the next value. |
| 18 | /// |
| 19 | /// Iterators in Rust don't express the ability to block, so this adapter |
| 20 | /// simply always calls `iter.next()` and returns that. |
| 21 | /// |
| 22 | /// ``` |
| 23 | /// # async fn dox() { |
| 24 | /// use tokio_stream::{self as stream, StreamExt}; |
| 25 | /// |
| 26 | /// let mut stream = stream::iter(vec![17, 19]); |
| 27 | /// |
| 28 | /// assert_eq!(stream.next().await, Some(17)); |
| 29 | /// assert_eq!(stream.next().await, Some(19)); |
| 30 | /// assert_eq!(stream.next().await, None); |
| 31 | /// # } |
| 32 | /// ``` |
| 33 | pub fn iter<I>(i: I) -> Iter<I::IntoIter> |
| 34 | where |
| 35 | I: IntoIterator, |
| 36 | { |
| 37 | Iter { |
| 38 | iter: i.into_iter(), |
| 39 | yield_amt: 0, |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | impl<I> Stream for Iter<I> |
| 44 | where |
| 45 | I: Iterator, |
| 46 | { |
| 47 | type Item = I::Item; |
| 48 | |
| 49 | fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<I::Item>> { |
| 50 | // TODO: add coop back |
| 51 | if self.yield_amt >= 32 { |
| 52 | self.yield_amt = 0; |
| 53 | |
| 54 | cx.waker().wake_by_ref(); |
| 55 | |
| 56 | Poll::Pending |
| 57 | } else { |
| 58 | self.yield_amt += 1; |
| 59 | |
| 60 | Poll::Ready(self.iter.next()) |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | fn size_hint(&self) -> (usize, Option<usize>) { |
| 65 | self.iter.size_hint() |
| 66 | } |
| 67 | } |
| 68 | |