1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

use core::mem::ManuallyDrop;
use core::pin::Pin;

use crate::stream::Stream;
use crate::task::{Context, Poll};

/// A stream that will repeatedly yield the same list of elements.
#[derive(Debug)]
pub struct Cycle<S> {
    orig: S,
    source: ManuallyDrop<S>,
}

impl<S> Cycle<S>
where
    S: Stream + Clone,
{
    pub(crate) fn new(source: S) -> Self {
        Self {
            orig: source.clone(),
            source: ManuallyDrop::new(source),
        }
    }
}

impl<S> Drop for Cycle<S> {
    fn drop(&mut self) {
        unsafe {
            ManuallyDrop::drop(&mut self.source);
        }
    }
}

impl<S> Stream for Cycle<S>
where
    S: Stream + Clone,
{
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        unsafe {
            let this = self.get_unchecked_mut();

            match futures_core::ready!(Pin::new_unchecked(&mut *this.source).poll_next(cx)) {
                Some(item) => Poll::Ready(Some(item)),
                None => {
                    ManuallyDrop::drop(&mut this.source);
                    this.source = ManuallyDrop::new(this.orig.clone());
                    Pin::new_unchecked(&mut *this.source).poll_next(cx)
                }
            }
        }
    }
}