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}
core/iter/sources/repeat_n.rs

core/iter/sources/
repeat_n.rs
