repeat_n.rs - source (original) (raw)
core/iter/sources/
repeat_n.rs
1use crate::fmt;
2use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
3use crate::mem::{self, MaybeUninit};
4use crate::num::NonZero;
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 let element = if count == 0 {
61 // `element` gets dropped eagerly.
62 MaybeUninit::uninit()
63 } else {
64 MaybeUninit::new(element)
65 };
66
67 RepeatN { element, count }
68}
69
70/// 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.
78 element: MaybeUninit<A>,
79}
80
81impl<A> RepeatN<A> {
82 /// Returns the element if it hasn't been dropped already.
83 fn element_ref(&self) -> Option<&A> {
84 if self.count > 0 {
85 // SAFETY: The count is non-zero, so it must be initialized.
86 Some(unsafe { self.element.assume_init_ref() })
87 } else {
88 None
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]
95 fn take_element(&mut self) -> Option<A> {
96 if self.count > 0 {
97 self.count = 0;
98 let 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.
101 unsafe { Some(element.assume_init()) }
102 } else {
103 None
104 }
105 }
106}
107
108#[stable(feature = "iter_repeat_n", since = "1.82.0")]
109impl<A: Clone> Clone for RepeatN<A> {
110 fn 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}
117
118#[stable(feature = "iter_repeat_n", since = "1.82.0")]
119impl<A: fmt::Debug> fmt::Debug for RepeatN<A> {
120 fn 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}
127
128#[stable(feature = "iter_repeat_n", since = "1.82.0")]
129impl<A> Drop for RepeatN<A> {
130 fn drop(&mut self) {
131 self.take_element();
132 }
133}
134
135#[stable(feature = "iter_repeat_n", since = "1.82.0")]
136impl<A: Clone> Iterator for RepeatN<A> {
137 type Item = A;
138
139 #[inline]
140 fn next(&mut self) -> Option<A> {
141 if self.count > 0 {
142 // SAFETY: Just checked it's not empty
143 unsafe { Some(self.next_unchecked()) }
144 } else {
145 None
146 }
147 }
148
149 #[inline]
150 fn size_hint(&self) -> (usize, Option<usize>) {
151 let len = self.len();
152 (len, Some(len))
153 }
154
155 #[inline]
156 fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
157 let len = self.count;
158
159 if skip >= len {
160 self.take_element();
161 }
162
163 if skip > len {
164 // SAFETY: we just checked that the difference is positive
165 Err(unsafe { NonZero::new_unchecked(skip - len) })
166 } else {
167 self.count = len - skip;
168 Ok(())
169 }
170 }
171
172 #[inline]
173 fn last(mut self) -> Option<A> {
174 self.take_element()
175 }
176
177 #[inline]
178 fn count(self) -> usize {
179 self.len()
180 }
181}
182
183#[stable(feature = "iter_repeat_n", since = "1.82.0")]
184impl<A: Clone> ExactSizeIterator for RepeatN<A> {
185 fn len(&self) -> usize {
186 self.count
187 }
188}
189
190#[stable(feature = "iter_repeat_n", since = "1.82.0")]
191impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
192 #[inline]
193 fn next_back(&mut self) -> Option<A> {
194 self.next()
195 }
196
197 #[inline]
198 fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
199 self.advance_by(n)
200 }
201
202 #[inline]
203 fn nth_back(&mut self, n: usize) -> Option<A> {
204 self.nth(n)
205 }
206}
207
208#[stable(feature = "iter_repeat_n", since = "1.82.0")]
209impl<A: Clone> FusedIterator for RepeatN<A> {}
210
211#[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]
216 unsafe fn next_unchecked(&mut self) -> Self::Item {
217 // SAFETY: The caller promised the iterator isn't empty
218 self.count = unsafe { self.count.unchecked_sub(1) };
219 if 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.
223 unsafe { 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.
226 let element = unsafe { self.element.assume_init_ref() };
227 A::clone(element)
228 }
229 }
230}