]>
Commit | Line | Data |
---|---|---|
487cf647 FG |
1 | use crate::iter::{FusedIterator, TrustedLen}; |
2 | use crate::mem::ManuallyDrop; | |
353b0b11 | 3 | use crate::num::NonZeroUsize; |
487cf647 FG |
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")] | |
62 | #[doc(hidden)] // waiting on ACP#120 to decide whether to expose publicly | |
63 | pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> { | |
64 | let mut element = ManuallyDrop::new(element); | |
65 | ||
66 | if count == 0 { | |
67 | // SAFETY: we definitely haven't dropped it yet, since we only just got | |
68 | // passed it in, and because the count is zero the instance we're about | |
69 | // to create won't drop it, so to avoid leaking we need to now. | |
70 | unsafe { ManuallyDrop::drop(&mut element) }; | |
71 | } | |
72 | ||
73 | RepeatN { element, count } | |
74 | } | |
75 | ||
76 | /// An iterator that repeats an element an exact number of times. | |
77 | /// | |
78 | /// This `struct` is created by the [`repeat_n()`] function. | |
79 | /// See its documentation for more. | |
80 | #[derive(Clone, Debug)] | |
81 | #[unstable(feature = "iter_repeat_n", issue = "104434")] | |
82 | #[doc(hidden)] // waiting on ACP#120 to decide whether to expose publicly | |
83 | pub struct RepeatN<A> { | |
84 | count: usize, | |
85 | // Invariant: has been dropped iff count == 0. | |
86 | element: ManuallyDrop<A>, | |
87 | } | |
88 | ||
89 | impl<A> RepeatN<A> { | |
90 | /// If we haven't already dropped the element, return it in an option. | |
91 | /// | |
92 | /// Clears the count so it won't be dropped again later. | |
93 | #[inline] | |
94 | fn take_element(&mut self) -> Option<A> { | |
95 | if self.count > 0 { | |
96 | self.count = 0; | |
97 | // SAFETY: We just set count to zero so it won't be dropped again, | |
98 | // and it used to be non-zero so it hasn't already been dropped. | |
99 | unsafe { Some(ManuallyDrop::take(&mut self.element)) } | |
100 | } else { | |
101 | None | |
102 | } | |
103 | } | |
104 | } | |
105 | ||
106 | #[unstable(feature = "iter_repeat_n", issue = "104434")] | |
107 | impl<A> Drop for RepeatN<A> { | |
108 | fn drop(&mut self) { | |
109 | self.take_element(); | |
110 | } | |
111 | } | |
112 | ||
113 | #[unstable(feature = "iter_repeat_n", issue = "104434")] | |
114 | impl<A: Clone> Iterator for RepeatN<A> { | |
115 | type Item = A; | |
116 | ||
117 | #[inline] | |
118 | fn next(&mut self) -> Option<A> { | |
119 | if self.count == 0 { | |
120 | return None; | |
121 | } | |
122 | ||
123 | self.count -= 1; | |
124 | Some(if self.count == 0 { | |
125 | // SAFETY: the check above ensured that the count used to be non-zero, | |
126 | // so element hasn't been dropped yet, and we just lowered the count to | |
127 | // zero so it won't be dropped later, and thus it's okay to take it here. | |
128 | unsafe { ManuallyDrop::take(&mut self.element) } | |
129 | } else { | |
9c376795 | 130 | A::clone(&self.element) |
487cf647 FG |
131 | }) |
132 | } | |
133 | ||
134 | #[inline] | |
135 | fn size_hint(&self) -> (usize, Option<usize>) { | |
136 | let len = self.len(); | |
137 | (len, Some(len)) | |
138 | } | |
139 | ||
140 | #[inline] | |
353b0b11 | 141 | fn advance_by(&mut self, skip: usize) -> Result<(), NonZeroUsize> { |
487cf647 FG |
142 | let len = self.count; |
143 | ||
144 | if skip >= len { | |
145 | self.take_element(); | |
146 | } | |
147 | ||
148 | if skip > len { | |
353b0b11 FG |
149 | // SAFETY: we just checked that the difference is positive |
150 | Err(unsafe { NonZeroUsize::new_unchecked(skip - len) }) | |
487cf647 FG |
151 | } else { |
152 | self.count = len - skip; | |
153 | Ok(()) | |
154 | } | |
155 | } | |
156 | ||
157 | #[inline] | |
158 | fn last(mut self) -> Option<A> { | |
159 | self.take_element() | |
160 | } | |
161 | ||
162 | #[inline] | |
163 | fn count(self) -> usize { | |
164 | self.len() | |
165 | } | |
166 | } | |
167 | ||
168 | #[unstable(feature = "iter_repeat_n", issue = "104434")] | |
169 | impl<A: Clone> ExactSizeIterator for RepeatN<A> { | |
170 | fn len(&self) -> usize { | |
171 | self.count | |
172 | } | |
173 | } | |
174 | ||
175 | #[unstable(feature = "iter_repeat_n", issue = "104434")] | |
176 | impl<A: Clone> DoubleEndedIterator for RepeatN<A> { | |
177 | #[inline] | |
178 | fn next_back(&mut self) -> Option<A> { | |
179 | self.next() | |
180 | } | |
181 | ||
182 | #[inline] | |
353b0b11 | 183 | fn advance_back_by(&mut self, n: usize) -> Result<(), NonZeroUsize> { |
487cf647 FG |
184 | self.advance_by(n) |
185 | } | |
186 | ||
187 | #[inline] | |
188 | fn nth_back(&mut self, n: usize) -> Option<A> { | |
189 | self.nth(n) | |
190 | } | |
191 | } | |
192 | ||
193 | #[unstable(feature = "iter_repeat_n", issue = "104434")] | |
194 | impl<A: Clone> FusedIterator for RepeatN<A> {} | |
195 | ||
196 | #[unstable(feature = "trusted_len", issue = "37572")] | |
197 | unsafe impl<A: Clone> TrustedLen for RepeatN<A> {} |