Repository: tokio-rs/tokio
File: tokio-stream/src/stream_map.rs (Rust)

Code before:
        &mut self,
        cx: &mut Context<'_>,
        buffer: &mut Vec<(K, V::Item)>,
        limit: usize,
    ) -> Poll<usize> {
        if limit == 0 || self.entries.is_empty() {
            return Poll::Ready(0);
        }

        let mut remaining = limit;
        let mut start = self::rand::thread_rng_n(self.entries.len() as u32) as usize;

        while remaining > 0 {
            match self.poll_one(cx, start) {
                Poll::Ready(Some((idx, val))) => {
                    remaining -= 1;
                    let key = self.entries[idx].0.clone();
                    buffer.push((key, val));

                    start = idx.wrapping_add(1) % self.entries.len();
                }
                Poll::Ready(None) | Poll::Pending => break,
            }
        }

        let added = limit - remaining;

        if added > 0 {
            Poll::Ready(added)
        } else if self.entries.is_empty() {
            Poll::Ready(0)
        } else {
            Poll::Pending
        }
    }
}

impl<K, V> Stream for StreamMap<K, V>
where
    K: Clone + Unpin,
    V: Stream + Unpin,
{
    type Item = (K, V::Item);

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if let Some((idx, val)) = ready!(self.poll_next_entry(cx)) {
            let key = self.entries[idx].0.clone();
            Poll::Ready(Some((key, val)))
        } else {
            Poll::Ready(None)
        }

Review comment (suggestion):
Seems simpler to just keep track of `added` from the beginning and change the while loop to `while added < limit`.

Code after:
        let start = self::rand::thread_rng_n(self.entries.len() as u32) as usize;
        let mut idx = start;

        while added < limit {
            // Indicates whether at least one stream returned a value when polled or not
            let mut should_loop = false;

            for _ in 0..self.entries.len() {
                let (_, stream) = &mut self.entries[idx];

                match Pin::new(stream).poll_next(cx) {
                    Poll::Ready(Some(val)) => {
                        added += 1;

                        let key = self.entries[idx].0.clone();
                        buffer.push((key, val));

                        should_loop = true;

                        idx = idx.wrapping_add(1) % self.entries.len();
                    }
                    Poll::Ready(None) => {
                        // Remove the entry
                        self.entries.swap_remove(idx);

                        // Check if this was the last entry, if so the cursor needs
                        // to wrap
                        if idx == self.entries.len() {
                            idx = 0;
                        } else if idx < start && start <= self.entries.len() {
                            // The stream being swapped into the current index has
                            // already been polled, so skip it.
                            idx = idx.wrapping_add(1) % self.entries.len();
                        }
                    }
                    Poll::Pending => {
                        idx = idx.wrapping_add(1) % self.entries.len();
                    }
                }
            }

            if !should_loop {
                break;
            }
        }

        if added > 0 {
            Poll::Ready(added)
        } else if self.entries.is_empty() {
            Poll::Ready(0)
        } else {