cycle_to_nth/
lib.rs

1//! A fork of the [`Cycle`] iterator adapter, to cycle until reaching the `n`th iteration.
2use core::iter::Iterator;
3
4#[derive(Clone, Debug)]
5#[must_use = "iterators are lazy and do nothing unless consumed"]
6pub struct CycleToNth<I> {
7    orig: I,
8    iter: I,
9    nth: usize,
10    n: usize,
11}
12
13impl<I: Clone> CycleToNth<I> {
14    pub fn new(iter: I, nth: usize) -> CycleToNth<I> {
15        CycleToNth {
16            orig: iter.clone(),
17            iter,
18            nth,
19            n: 0,
20        }
21    }
22}
23
24impl<I> Iterator for CycleToNth<I>
25where
26    I: Clone + Iterator,
27{
28    type Item = <I as Iterator>::Item;
29
30    #[inline]
31    fn next(&mut self) -> Option<<I as Iterator>::Item> {
32        if self.n == self.nth {
33            // Terminate on the nth iteration
34            None
35        } else {
36            self.n += 1;
37            match self.iter.next() {
38                None => {
39                    self.iter = self.orig.clone();
40                    self.iter.next()
41                }
42                y => y,
43            }
44        }
45    }
46
47    #[inline]
48    fn size_hint(&self) -> (usize, Option<usize>) {
49        // the CycleToNth iterator is either empty of infinite
50        match self.orig.size_hint() {
51            sz @ (0, Some(0)) => sz,
52            (0, _) => (0, None),
53            _ => (usize::MAX, None),
54        }
55    }
56}
57
58pub trait CycleToNthAdapter
59where
60    Self: Sized + Clone,
61{
62    fn cycle_to_nth(self, nth: usize) -> CycleToNth<Self>;
63}
64
65impl<I> CycleToNthAdapter for I
66where
67    Self: Sized + Clone,
68    I: Iterator,
69{
70    fn cycle_to_nth(self, nth: usize) -> CycleToNth<Self> {
71        CycleToNth::new(self, nth)
72    }
73}
74
75#[cfg(test)]
76mod test {
77    use crate::*;
78
79    #[test]
80    fn extend_with_cycles() {
81        let v = vec![1u32];
82        assert_eq!(v.iter().cycle_to_nth(5).sum::<u32>(), 5u32);
83    }
84
85    #[test]
86    fn truncate_to_nth() {
87        let v = vec![1u32; 10];
88        assert_eq!(v.iter().cycle_to_nth(5).sum::<u32>(), 5u32);
89    }
90}
cycle_to_nth/lib.rs

cycle_to_nth/
lib.rs
