1use crate::iter::{FusedIterator, TrustedLen};
2use crate::mem::ManuallyDrop;
3use crate::num::NonZero;
4
5/// Creates a new iterator that repeats a single element a given number of times.
6///
7/// The `repeat_n()` function repeats a single value exactly `n` times.
8///
9/// This is very similar to using [`repeat()`] with [`Iterator::take()`],
10/// but there are two differences:
11/// - `repeat_n()` can return the original value, rather than always cloning.
12/// - `repeat_n()` produces an [`ExactSizeIterator`].
13///
14/// [`repeat()`]: crate::iter::repeat
15///
16/// # Examples
17///
18/// Basic usage:
19///
20/// ```
21/// #![feature(iter_repeat_n)]
22/// use std::iter;
23///
24/// // four of the number four:
25/// let mut four_fours = iter::repeat_n(4, 4);
26///
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/// assert_eq!(Some(4), four_fours.next());
31///
32/// // no more fours
33/// assert_eq!(None, four_fours.next());
34/// ```
35///
36/// For non-`Copy` types,
37///
38/// ```
39/// #![feature(iter_repeat_n)]
40/// use std::iter;
41///
42/// let v: Vec<i32> = Vec::with_capacity(123);
43/// let mut it = iter::repeat_n(v, 5);
44///
45/// for i in 0..4 {
46/// // It starts by cloning things
47/// let cloned = it.next().unwrap();
48/// assert_eq!(cloned.len(), 0);
49/// assert_eq!(cloned.capacity(), 0);
50/// }
51///
52/// // ... but the last item is the original one
53/// let last = it.next().unwrap();
54/// assert_eq!(last.len(), 0);
55/// assert_eq!(last.capacity(), 123);
56///
57/// // ... and now we're done
58/// assert_eq!(None, it.next());
59/// ```
60#[inline]
61#[unstable(feature = "iter_repeat_n", issue = "104434")]
62pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
63 let mut element: ManuallyDrop = ManuallyDrop::new(element);
64
65 if count == 0 {
66 // SAFETY: we definitely haven't dropped it yet, since we only just got
67 // passed it in, and because the count is zero the instance we're about
68 // to create won't drop it, so to avoid leaking we need to now.
69 unsafe { ManuallyDrop::drop(&mut element) };
70 }
71
72 RepeatN { element, count }
73}
74
75/// An iterator that repeats an element an exact number of times.
76///
77/// This `struct` is created by the [`repeat_n()`] function.
78/// See its documentation for more.
79#[derive(Clone, Debug)]
80#[unstable(feature = "iter_repeat_n", issue = "104434")]
81pub struct RepeatN<A> {
82 count: usize,
83 // Invariant: has been dropped iff count == 0.
84 element: ManuallyDrop<A>,
85}
86
87impl<A> RepeatN<A> {
88 /// If we haven't already dropped the element, return it in an option.
89 ///
90 /// Clears the count so it won't be dropped again later.
91 #[inline]
92 fn take_element(&mut self) -> Option<A> {
93 if self.count > 0 {
94 self.count = 0;
95 // SAFETY: We just set count to zero so it won't be dropped again,
96 // and it used to be non-zero so it hasn't already been dropped.
97 unsafe { Some(ManuallyDrop::take(&mut self.element)) }
98 } else {
99 None
100 }
101 }
102}
103
104#[unstable(feature = "iter_repeat_n", issue = "104434")]
105impl<A> Drop for RepeatN<A> {
106 fn drop(&mut self) {
107 self.take_element();
108 }
109}
110
111#[unstable(feature = "iter_repeat_n", issue = "104434")]
112impl<A: Clone> Iterator for RepeatN<A> {
113 type Item = A;
114
115 #[inline]
116 fn next(&mut self) -> Option<A> {
117 if self.count == 0 {
118 return None;
119 }
120
121 self.count -= 1;
122 Some(if self.count == 0 {
123 // SAFETY: the check above ensured that the count used to be non-zero,
124 // so element hasn't been dropped yet, and we just lowered the count to
125 // zero so it won't be dropped later, and thus it's okay to take it here.
126 unsafe { ManuallyDrop::take(&mut self.element) }
127 } else {
128 A::clone(&self.element)
129 })
130 }
131
132 #[inline]
133 fn size_hint(&self) -> (usize, Option<usize>) {
134 let len = self.len();
135 (len, Some(len))
136 }
137
138 #[inline]
139 fn advance_by(&mut self, skip: usize) -> Result<(), NonZero<usize>> {
140 let len = self.count;
141
142 if skip >= len {
143 self.take_element();
144 }
145
146 if skip > len {
147 // SAFETY: we just checked that the difference is positive
148 Err(unsafe { NonZero::new_unchecked(skip - len) })
149 } else {
150 self.count = len - skip;
151 Ok(())
152 }
153 }
154
155 #[inline]
156 fn last(mut self) -> Option<A> {
157 self.take_element()
158 }
159
160 #[inline]
161 fn count(self) -> usize {
162 self.len()
163 }
164}
165
166#[unstable(feature = "iter_repeat_n", issue = "104434")]
167impl<A: Clone> ExactSizeIterator for RepeatN<A> {
168 fn len(&self) -> usize {
169 self.count
170 }
171}
172
173#[unstable(feature = "iter_repeat_n", issue = "104434")]
174impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
175 #[inline]
176 fn next_back(&mut self) -> Option<A> {
177 self.next()
178 }
179
180 #[inline]
181 fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
182 self.advance_by(n)
183 }
184
185 #[inline]
186 fn nth_back(&mut self, n: usize) -> Option<A> {
187 self.nth(n)
188 }
189}
190
191#[unstable(feature = "iter_repeat_n", issue = "104434")]
192impl<A: Clone> FusedIterator for RepeatN<A> {}
193
194#[unstable(feature = "trusted_len", issue = "37572")]
195unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
196