1use crate::fmt;
2use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
3use crate::mem::{self, MaybeUninit};
4use crate::num::NonZero;
56/// 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> {
60let element = if count == 0 {
61// `element` gets dropped eagerly.
62MaybeUninit::uninit()
63 } else {
64 MaybeUninit::new(element)
65 };
6667 RepeatN { element, count }
68}
6970/// An iterator that repeats an element an exact number of times.
71///
72/// This `struct` is created by the [`repeat_n()`] function.
73/// See its documentation for more.
74#[stable(feature = "iter_repeat_n", since = "1.82.0")]
75pub struct RepeatN<A> {
76 count: usize,
77// Invariant: uninit iff count == 0.
78element: MaybeUninit<A>,
79}
8081impl<A> RepeatN<A> {
82/// Returns the element if it hasn't been dropped already.
83fn element_ref(&self) -> Option<&A> {
84if self.count > 0 {
85// SAFETY: The count is non-zero, so it must be initialized.
86Some(unsafe { self.element.assume_init_ref() })
87 } else {
88None
89}
90 }
91/// If we haven't already dropped the element, return it in an option.
92 ///
93 /// Clears the count so it won't be dropped again later.
94#[inline]
95fn take_element(&mut self) -> Option<A> {
96if self.count > 0 {
97self.count = 0;
98let element = mem::replace(&mut self.element, MaybeUninit::uninit());
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.
101unsafe { Some(element.assume_init()) }
102 } else {
103None
104}
105 }
106}
107108#[stable(feature = "iter_repeat_n", since = "1.82.0")]
109impl<A: Clone> Clone for RepeatN<A> {
110fn clone(&self) -> RepeatN<A> {
111 RepeatN {
112 count: self.count,
113 element: self.element_ref().cloned().map_or_else(MaybeUninit::uninit, MaybeUninit::new),
114 }
115 }
116}
117118#[stable(feature = "iter_repeat_n", since = "1.82.0")]
119impl<A: fmt::Debug> fmt::Debug for RepeatN<A> {
120fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121 f.debug_struct("RepeatN")
122 .field("count", &self.count)
123 .field("element", &self.element_ref())
124 .finish()
125 }
126}
127128#[stable(feature = "iter_repeat_n", since = "1.82.0")]
129impl<A> Drop for RepeatN<A> {
130fn drop(&mut self) {
131self.take_element();
132 }
133}
134135#[stable(feature = "iter_repeat_n", since = "1.82.0")]
136impl<A: Clone> Iterator for RepeatN<A> {
137type Item = A;
138139#[inline]
140fn next(&mut self) -> Option<A> {
141if self.count > 0 {
142// SAFETY: Just checked it's not empty
143unsafe { Some(self.next_unchecked()) }
144 } else {
145None
146}
147 }
148149#[inline]
150fn size_hint(&self) -> (usize, Option<usize>) {
151let len = self.len();
152 (len, Some(len))
153 }
154155#[inline]
156fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
157let len = self.count;
158159if skip >= len {
160self.take_element();
161 }
162163if skip > len {
164// SAFETY: we just checked that the difference is positive
165Err(unsafe { NonZero::new_unchecked(skip - len) })
166 } else {
167self.count = len - skip;
168Ok(())
169 }
170 }
171172#[inline]
173fn last(mut self) -> Option<A> {
174self.take_element()
175 }
176177#[inline]
178fn count(self) -> usize {
179self.len()
180 }
181}
182183#[stable(feature = "iter_repeat_n", since = "1.82.0")]
184impl<A: Clone> ExactSizeIterator for RepeatN<A> {
185fn len(&self) -> usize {
186self.count
187 }
188}
189190#[stable(feature = "iter_repeat_n", since = "1.82.0")]
191impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
192#[inline]
193fn next_back(&mut self) -> Option<A> {
194self.next()
195 }
196197#[inline]
198fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
199self.advance_by(n)
200 }
201202#[inline]
203fn nth_back(&mut self, n: usize) -> Option<A> {
204self.nth(n)
205 }
206}
207208#[stable(feature = "iter_repeat_n", since = "1.82.0")]
209impl<A: Clone> FusedIterator for RepeatN<A> {}
210211#[unstable(feature = "trusted_len", issue = "37572")]
212unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
213#[stable(feature = "iter_repeat_n", since = "1.82.0")]
214impl<A: Clone> UncheckedIterator for RepeatN<A> {
215#[inline]
216unsafe fn next_unchecked(&mut self) -> Self::Item {
217// SAFETY: The caller promised the iterator isn't empty
218self.count = unsafe { self.count.unchecked_sub(1) };
219if self.count == 0 {
220// SAFETY: the check above ensured that the count used to be non-zero,
221 // so element hasn't been dropped yet, and we just lowered the count to
222 // zero so it won't be dropped later, and thus it's okay to take it here.
223unsafe { mem::replace(&mut self.element, MaybeUninit::uninit()).assume_init() }
224 } else {
225// SAFETY: the count is non-zero, so it must have not been dropped yet.
226let element = unsafe { self.element.assume_init_ref() };
227 A::clone(element)
228 }
229 }
230}