core/iter/sources/
repeat_n.rs1use crate::fmt;
2use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
3use crate::mem::MaybeUninit;
4use crate::num::NonZero;
5use crate::ops::{NeverShortCircuit, Try};
6
7#[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 = if count == 0 {
62        MaybeUninit::uninit()
64    } else {
65        MaybeUninit::new(element)
66    };
67
68    RepeatN { element, count }
69}
70
71#[stable(feature = "iter_repeat_n", since = "1.82.0")]
76pub struct RepeatN<A> {
77    count: usize,
78    element: MaybeUninit<A>,
80}
81
82impl<A> RepeatN<A> {
83    fn element_ref(&self) -> Option<&A> {
85        if self.count > 0 {
86            Some(unsafe { self.element.assume_init_ref() })
88        } else {
89            None
90        }
91    }
92    #[inline]
96    fn take_element(&mut self) -> Option<A> {
97        if self.count > 0 {
98            self.count = 0;
99            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(MaybeUninit::uninit, 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.debug_struct("RepeatN")
123            .field("count", &self.count)
124            .field("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            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            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                acc = f(acc, unsafe { self.element.assume_init_ref().clone() })?;
184            }
185
186            self.count -= 1;
190            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        self.count = unsafe { self.count.unchecked_sub(1) };
270        if self.count == 0 {
271            unsafe { self.element.assume_init_read() }
275        } else {
276            let element = unsafe { self.element.assume_init_ref() };
278            A::clone(element)
279        }
280    }
281}