core/iter/sources/
repeat_n.rs

1use crate::fmt;
2use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
3use crate::num::NonZero;
4use crate::ops::Try;
5
6/// Creates a new iterator that repeats a single element a given number of times.
7///
8/// The `repeat_n()` function repeats a single value exactly `n` times.
9///
10/// This is very similar to using [`repeat()`] with [`Iterator::take()`],
11/// but `repeat_n()` can return the original value, rather than always cloning.
12///
13/// [`repeat()`]: crate::iter::repeat
14///
15/// # Examples
16///
17/// Basic usage:
18///
19/// ```
20/// use std::iter;
21///
22/// // four of the number four:
23/// let mut four_fours = iter::repeat_n(4, 4);
24///
25/// assert_eq!(Some(4), four_fours.next());
26/// assert_eq!(Some(4), four_fours.next());
27/// assert_eq!(Some(4), four_fours.next());
28/// assert_eq!(Some(4), four_fours.next());
29///
30/// // no more fours
31/// assert_eq!(None, four_fours.next());
32/// ```
33///
34/// For non-`Copy` types,
35///
36/// ```
37/// use std::iter;
38///
39/// let v: Vec<i32> = Vec::with_capacity(123);
40/// let mut it = iter::repeat_n(v, 5);
41///
42/// for i in 0..4 {
43///     // It starts by cloning things
44///     let cloned = it.next().unwrap();
45///     assert_eq!(cloned.len(), 0);
46///     assert_eq!(cloned.capacity(), 0);
47/// }
48///
49/// // ... but the last item is the original one
50/// let last = it.next().unwrap();
51/// assert_eq!(last.len(), 0);
52/// assert_eq!(last.capacity(), 123);
53///
54/// // ... and now we're done
55/// assert_eq!(None, it.next());
56/// ```
57#[inline]
58#[stable(feature = "iter_repeat_n", since = "1.82.0")]
59pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
60    RepeatN { inner: RepeatNInner::new(element, count) }
61}
62
63#[derive(Clone, Copy)]
64struct RepeatNInner<T> {
65    count: NonZero<usize>,
66    element: T,
67}
68
69impl<T> RepeatNInner<T> {
70    fn new(element: T, count: usize) -> Option<Self> {
71        let count = NonZero::<usize>::new(count)?;
72        Some(Self { element, count })
73    }
74}
75
76/// An iterator that repeats an element an exact number of times.
77///
78/// This `struct` is created by the [`repeat_n()`] function.
79/// See its documentation for more.
80#[stable(feature = "iter_repeat_n", since = "1.82.0")]
81#[derive(Clone)]
82pub struct RepeatN<A> {
83    inner: Option<RepeatNInner<A>>,
84}
85
86impl<A> RepeatN<A> {
87    /// If we haven't already dropped the element, return it in an option.
88    #[inline]
89    fn take_element(&mut self) -> Option<A> {
90        self.inner.take().map(|inner| inner.element)
91    }
92}
93
94#[stable(feature = "iter_repeat_n", since = "1.82.0")]
95impl<A: fmt::Debug> fmt::Debug for RepeatN<A> {
96    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97        let (count, element) = match self.inner.as_ref() {
98            Some(inner) => (inner.count.get(), Some(&inner.element)),
99            None => (0, None),
100        };
101        f.debug_struct("RepeatN").field("count", &count).field("element", &element).finish()
102    }
103}
104
105#[stable(feature = "iter_repeat_n", since = "1.82.0")]
106impl<A: Clone> Iterator for RepeatN<A> {
107    type Item = A;
108
109    #[inline]
110    fn next(&mut self) -> Option<A> {
111        let inner = self.inner.as_mut()?;
112        let count = inner.count.get();
113
114        if let Some(decremented) = NonZero::<usize>::new(count - 1) {
115            // Order of these is important for optimization
116            let tmp = inner.element.clone();
117            inner.count = decremented;
118            return Some(tmp);
119        }
120
121        return self.take_element();
122    }
123
124    #[inline]
125    fn size_hint(&self) -> (usize, Option<usize>) {
126        let len = self.len();
127        (len, Some(len))
128    }
129
130    #[inline]
131    fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
132        let Some(inner) = self.inner.as_mut() else {
133            return NonZero::<usize>::new(skip).map(Err).unwrap_or(Ok(()));
134        };
135
136        let len = inner.count.get();
137
138        if let Some(new_len) = len.checked_sub(skip).and_then(NonZero::<usize>::new) {
139            inner.count = new_len;
140            return Ok(());
141        }
142
143        self.inner = None;
144        return NonZero::<usize>::new(skip - len).map(Err).unwrap_or(Ok(()));
145    }
146
147    #[inline]
148    fn last(mut self) -> Option<A> {
149        self.take_element()
150    }
151
152    #[inline]
153    fn count(self) -> usize {
154        self.len()
155    }
156}
157
158#[stable(feature = "iter_repeat_n", since = "1.82.0")]
159impl<A: Clone> ExactSizeIterator for RepeatN<A> {
160    fn len(&self) -> usize {
161        self.inner.as_ref().map(|inner| inner.count.get()).unwrap_or(0)
162    }
163}
164
165#[stable(feature = "iter_repeat_n", since = "1.82.0")]
166impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
167    #[inline]
168    fn next_back(&mut self) -> Option<A> {
169        self.next()
170    }
171
172    #[inline]
173    fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
174        self.advance_by(n)
175    }
176
177    #[inline]
178    fn nth_back(&mut self, n: usize) -> Option<A> {
179        self.nth(n)
180    }
181
182    #[inline]
183    fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
184    where
185        F: FnMut(B, A) -> R,
186        R: Try<Output = B>,
187    {
188        self.try_fold(init, f)
189    }
190
191    #[inline]
192    fn rfold<B, F>(self, init: B, f: F) -> B
193    where
194        F: FnMut(B, A) -> B,
195    {
196        self.fold(init, f)
197    }
198}
199
200#[stable(feature = "iter_repeat_n", since = "1.82.0")]
201impl<A: Clone> FusedIterator for RepeatN<A> {}
202
203#[unstable(feature = "trusted_len", issue = "37572")]
204unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
205#[stable(feature = "iter_repeat_n", since = "1.82.0")]
206impl<A: Clone> UncheckedIterator for RepeatN<A> {}
core/iter/sources/repeat_n.rs

core/iter/sources/
repeat_n.rs
