itertools/
repeatn.rs

1use std::iter::FusedIterator;
2
3/// An iterator that produces *n* repetitions of an element.
4///
5/// See [`repeat_n()`](crate::repeat_n) for more information.
6#[must_use = "iterators are lazy and do nothing unless consumed"]
7#[derive(Clone, Debug)]
8pub struct RepeatN<A> {
9    pub(crate) elt: Option<A>,
10    n: usize,
11}
12
13/// Create an iterator that produces `n` repetitions of `element`.
14pub fn repeat_n<A>(element: A, n: usize) -> RepeatN<A>
15where
16    A: Clone,
17{
18    if n == 0 {
19        RepeatN { elt: None, n }
20    } else {
21        RepeatN {
22            elt: Some(element),
23            n,
24        }
25    }
26}
27
28impl<A> Iterator for RepeatN<A>
29where
30    A: Clone,
31{
32    type Item = A;
33
34    fn next(&mut self) -> Option<Self::Item> {
35        if self.n > 1 {
36            self.n -= 1;
37            self.elt.as_ref().cloned()
38        } else {
39            self.n = 0;
40            self.elt.take()
41        }
42    }
43
44    fn size_hint(&self) -> (usize, Option<usize>) {
45        (self.n, Some(self.n))
46    }
47
48    fn fold<B, F>(self, mut init: B, mut f: F) -> B
49    where
50        F: FnMut(B, Self::Item) -> B,
51    {
52        match self {
53            Self { elt: Some(elt), n } => {
54                debug_assert!(n > 0);
55                init = (1..n).map(|_| elt.clone()).fold(init, &mut f);
56                f(init, elt)
57            }
58            _ => init,
59        }
60    }
61}
62
63impl<A> DoubleEndedIterator for RepeatN<A>
64where
65    A: Clone,
66{
67    #[inline]
68    fn next_back(&mut self) -> Option<Self::Item> {
69        self.next()
70    }
71
72    #[inline]
73    fn rfold<B, F>(self, init: B, f: F) -> B
74    where
75        F: FnMut(B, Self::Item) -> B,
76    {
77        self.fold(init, f)
78    }
79}
80
81impl<A> ExactSizeIterator for RepeatN<A> where A: Clone {}
82
83impl<A> FusedIterator for RepeatN<A> where A: Clone {}