1use crate::fmt;
2use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
3use crate::mem::MaybeUninit;
4use crate::num::NonZero;
5use crate::ops::{NeverShortCircuit, Try};
6
7/// Creates a new iterator that repeats a single element a given number of times.
8///
9/// The `repeat_n()` function repeats a single value exactly `n` times.
10///
11/// This is very similar to using [`repeat()`] with [`Iterator::take()`],
12/// but `repeat_n()` can return the original value, rather than always cloning.
13///
14/// [`repeat()`]: crate::iter::repeat
15///
16/// # Examples
17///
18/// Basic usage:
19///
20/// ```
21/// use std::iter;
22///
23/// // four of the number four:
24/// let mut four_fours = iter::repeat_n(4, 4);
25///
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/// assert_eq!(Some(4), four_fours.next());
30///
31/// // no more fours
32/// assert_eq!(None, four_fours.next());
33/// ```
34///
35/// For non-`Copy` types,
36///
37/// ```
38/// use std::iter;
39///
40/// let v: Vec<i32> = Vec::with_capacity(123);
41/// let mut it = iter::repeat_n(v, 5);
42///
43/// for i in 0..4 {
44/// // It starts by cloning things
45/// let cloned = it.next().unwrap();
46/// assert_eq!(cloned.len(), 0);
47/// assert_eq!(cloned.capacity(), 0);
48/// }
49///
50/// // ... but the last item is the original one
51/// let last = it.next().unwrap();
52/// assert_eq!(last.len(), 0);
53/// assert_eq!(last.capacity(), 123);
54///
55/// // ... and now we're done
56/// assert_eq!(None, it.next());
57/// ```
58#[inline]
59#[stable(feature = "iter_repeat_n", since = "1.82.0")]
60pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
61 let element: MaybeUninit = if count == 0 {
62 // `element` gets dropped eagerly.
63 MaybeUninit::uninit()
64 } else {
65 MaybeUninit::new(val:element)
66 };
67
68 RepeatN { element, count }
69}
70
71/// An iterator that repeats an element an exact number of times.
72///
73/// This `struct` is created by the [`repeat_n()`] function.
74/// See its documentation for more.
75#[stable(feature = "iter_repeat_n", since = "1.82.0")]
76pub struct RepeatN<A> {
77 count: usize,
78 // Invariant: uninit iff count == 0.
79 element: MaybeUninit<A>,
80}
81
82impl<A> RepeatN<A> {
83 /// Returns the element if it hasn't been dropped already.
84 fn element_ref(&self) -> Option<&A> {
85 if self.count > 0 {
86 // SAFETY: The count is non-zero, so it must be initialized.
87 Some(unsafe { self.element.assume_init_ref() })
88 } else {
89 None
90 }
91 }
92 /// If we haven't already dropped the element, return it in an option.
93 ///
94 /// Clears the count so it won't be dropped again later.
95 #[inline]
96 fn take_element(&mut self) -> Option<A> {
97 if self.count > 0 {
98 self.count = 0;
99 // SAFETY: We just set count to zero so it won't be dropped again,
100 // and it used to be non-zero so it hasn't already been dropped.
101 let element = unsafe { self.element.assume_init_read() };
102 Some(element)
103 } else {
104 None
105 }
106 }
107}
108
109#[stable(feature = "iter_repeat_n", since = "1.82.0")]
110impl<A: Clone> Clone for RepeatN<A> {
111 fn clone(&self) -> RepeatN<A> {
112 RepeatN {
113 count: self.count,
114 element: self.element_ref().cloned().map_or_else(default:MaybeUninit::uninit, f:MaybeUninit::new),
115 }
116 }
117}
118
119#[stable(feature = "iter_repeat_n", since = "1.82.0")]
120impl<A: fmt::Debug> fmt::Debug for RepeatN<A> {
121 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122 f&mut DebugStruct<'_, '_>.debug_struct("RepeatN")
123 .field("count", &self.count)
124 .field(name:"element", &self.element_ref())
125 .finish()
126 }
127}
128
129#[stable(feature = "iter_repeat_n", since = "1.82.0")]
130impl<A> Drop for RepeatN<A> {
131 fn drop(&mut self) {
132 self.take_element();
133 }
134}
135
136#[stable(feature = "iter_repeat_n", since = "1.82.0")]
137impl<A: Clone> Iterator for RepeatN<A> {
138 type Item = A;
139
140 #[inline]
141 fn next(&mut self) -> Option<A> {
142 if self.count > 0 {
143 // SAFETY: Just checked it's not empty
144 unsafe { Some(self.next_unchecked()) }
145 } else {
146 None
147 }
148 }
149
150 #[inline]
151 fn size_hint(&self) -> (usize, Option<usize>) {
152 let len = self.len();
153 (len, Some(len))
154 }
155
156 #[inline]
157 fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
158 let len = self.count;
159
160 if skip >= len {
161 self.take_element();
162 }
163
164 if skip > len {
165 // SAFETY: we just checked that the difference is positive
166 Err(unsafe { NonZero::new_unchecked(skip - len) })
167 } else {
168 self.count = len - skip;
169 Ok(())
170 }
171 }
172
173 fn try_fold<B, F, R>(&mut self, mut acc: B, mut f: F) -> R
174 where
175 F: FnMut(B, A) -> R,
176 R: Try<Output = B>,
177 {
178 if self.count > 0 {
179 while self.count > 1 {
180 self.count -= 1;
181 // SAFETY: the count was larger than 1, so the element is
182 // initialized and hasn't been dropped.
183 acc = f(acc, unsafe { self.element.assume_init_ref().clone() })?;
184 }
185
186 // We could just set the count to zero directly, but doing it this
187 // way should make it easier for the optimizer to fold this tail
188 // into the loop when `clone()` is equivalent to copying.
189 self.count -= 1;
190 // SAFETY: we just set the count to zero from one, so the element
191 // is still initialized, has not been dropped yet and will not be
192 // accessed by future calls.
193 f(acc, unsafe { self.element.assume_init_read() })
194 } else {
195 try { acc }
196 }
197 }
198
199 fn fold<B, F>(mut self, init: B, f: F) -> B
200 where
201 F: FnMut(B, A) -> B,
202 {
203 self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0
204 }
205
206 #[inline]
207 fn last(mut self) -> Option<A> {
208 self.take_element()
209 }
210
211 #[inline]
212 fn count(self) -> usize {
213 self.len()
214 }
215}
216
217#[stable(feature = "iter_repeat_n", since = "1.82.0")]
218impl<A: Clone> ExactSizeIterator for RepeatN<A> {
219 fn len(&self) -> usize {
220 self.count
221 }
222}
223
224#[stable(feature = "iter_repeat_n", since = "1.82.0")]
225impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
226 #[inline]
227 fn next_back(&mut self) -> Option<A> {
228 self.next()
229 }
230
231 #[inline]
232 fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
233 self.advance_by(n)
234 }
235
236 #[inline]
237 fn nth_back(&mut self, n: usize) -> Option<A> {
238 self.nth(n)
239 }
240
241 #[inline]
242 fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
243 where
244 F: FnMut(B, A) -> R,
245 R: Try<Output = B>,
246 {
247 self.try_fold(init, f)
248 }
249
250 #[inline]
251 fn rfold<B, F>(self, init: B, f: F) -> B
252 where
253 F: FnMut(B, A) -> B,
254 {
255 self.fold(init, f)
256 }
257}
258
259#[stable(feature = "iter_repeat_n", since = "1.82.0")]
260impl<A: Clone> FusedIterator for RepeatN<A> {}
261
262#[unstable(feature = "trusted_len", issue = "37572")]
263unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
264#[stable(feature = "iter_repeat_n", since = "1.82.0")]
265impl<A: Clone> UncheckedIterator for RepeatN<A> {
266 #[inline]
267 unsafe fn next_unchecked(&mut self) -> Self::Item {
268 // SAFETY: The caller promised the iterator isn't empty
269 self.count = unsafe { self.count.unchecked_sub(1) };
270 if self.count == 0 {
271 // SAFETY: the check above ensured that the count used to be non-zero,
272 // so element hasn't been dropped yet, and we just lowered the count to
273 // zero so it won't be dropped later, and thus it's okay to take it here.
274 unsafe { self.element.assume_init_read() }
275 } else {
276 // SAFETY: the count is non-zero, so it must have not been dropped yet.
277 let element: &A = unsafe { self.element.assume_init_ref() };
278 A::clone(self:element)
279 }
280 }
281}
282