]> git.lizzy.rs Git - rust.git/blob - library/core/src/array/iter.rs
implement fold() on array::IntoIter to improve flatten().collect() perf
[rust.git] / library / core / src / array / iter.rs
1 //! Defines the `IntoIter` owned iterator for arrays.
2
3 use crate::{
4     fmt,
5     iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
6     mem::{self, MaybeUninit},
7     ops::Range,
8     ptr,
9 };
10
11 /// A by-value [array] iterator.
12 #[stable(feature = "array_value_iter", since = "1.51.0")]
13 pub struct IntoIter<T, const N: usize> {
14     /// This is the array we are iterating over.
15     ///
16     /// Elements with index `i` where `alive.start <= i < alive.end` have not
17     /// been yielded yet and are valid array entries. Elements with indices `i
18     /// < alive.start` or `i >= alive.end` have been yielded already and must
19     /// not be accessed anymore! Those dead elements might even be in a
20     /// completely uninitialized state!
21     ///
22     /// So the invariants are:
23     /// - `data[alive]` is alive (i.e. contains valid elements)
24     /// - `data[..alive.start]` and `data[alive.end..]` are dead (i.e. the
25     ///   elements were already read and must not be touched anymore!)
26     data: [MaybeUninit<T>; N],
27
28     /// The elements in `data` that have not been yielded yet.
29     ///
30     /// Invariants:
31     /// - `alive.start <= alive.end`
32     /// - `alive.end <= N`
33     alive: Range<usize>,
34 }
35
36 impl<T, const N: usize> IntoIter<T, N> {
37     /// Creates a new iterator over the given `array`.
38     ///
39     /// *Note*: this method might be deprecated in the future,
40     /// after [`IntoIterator` is implemented for arrays][array-into-iter].
41     ///
42     /// # Examples
43     ///
44     /// ```
45     /// use std::array;
46     ///
47     /// for value in array::IntoIter::new([1, 2, 3, 4, 5]) {
48     ///     // The type of `value` is a `i32` here, instead of `&i32`
49     ///     let _: i32 = value;
50     /// }
51     /// ```
52     /// [array-into-iter]: https://github.com/rust-lang/rust/pull/65819
53     #[stable(feature = "array_value_iter", since = "1.51.0")]
54     pub fn new(array: [T; N]) -> Self {
55         // SAFETY: The transmute here is actually safe. The docs of `MaybeUninit`
56         // promise:
57         //
58         // > `MaybeUninit<T>` is guaranteed to have the same size and alignment
59         // > as `T`.
60         //
61         // The docs even show a transmute from an array of `MaybeUninit<T>` to
62         // an array of `T`.
63         //
64         // With that, this initialization satisfies the invariants.
65
66         // FIXME(LukasKalbertodt): actually use `mem::transmute` here, once it
67         // works with const generics:
68         //     `mem::transmute::<[T; N], [MaybeUninit<T>; N]>(array)`
69         //
70         // Until then, we can use `mem::transmute_copy` to create a bitwise copy
71         // as a different type, then forget `array` so that it is not dropped.
72         unsafe {
73             let iter = Self { data: mem::transmute_copy(&array), alive: 0..N };
74             mem::forget(array);
75             iter
76         }
77     }
78
79     /// Returns an immutable slice of all elements that have not been yielded
80     /// yet.
81     #[stable(feature = "array_value_iter", since = "1.51.0")]
82     pub fn as_slice(&self) -> &[T] {
83         // SAFETY: We know that all elements within `alive` are properly initialized.
84         unsafe {
85             let slice = self.data.get_unchecked(self.alive.clone());
86             MaybeUninit::slice_assume_init_ref(slice)
87         }
88     }
89
90     /// Returns a mutable slice of all elements that have not been yielded yet.
91     #[stable(feature = "array_value_iter", since = "1.51.0")]
92     pub fn as_mut_slice(&mut self) -> &mut [T] {
93         // SAFETY: We know that all elements within `alive` are properly initialized.
94         unsafe {
95             let slice = self.data.get_unchecked_mut(self.alive.clone());
96             MaybeUninit::slice_assume_init_mut(slice)
97         }
98     }
99 }
100
101 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
102 impl<T, const N: usize> Iterator for IntoIter<T, N> {
103     type Item = T;
104     fn next(&mut self) -> Option<Self::Item> {
105         // Get the next index from the front.
106         //
107         // Increasing `alive.start` by 1 maintains the invariant regarding
108         // `alive`. However, due to this change, for a short time, the alive
109         // zone is not `data[alive]` anymore, but `data[idx..alive.end]`.
110         self.alive.next().map(|idx| {
111             // Read the element from the array.
112             // SAFETY: `idx` is an index into the former "alive" region of the
113             // array. Reading this element means that `data[idx]` is regarded as
114             // dead now (i.e. do not touch). As `idx` was the start of the
115             // alive-zone, the alive zone is now `data[alive]` again, restoring
116             // all invariants.
117             unsafe { self.data.get_unchecked(idx).assume_init_read() }
118         })
119     }
120
121     fn size_hint(&self) -> (usize, Option<usize>) {
122         let len = self.len();
123         (len, Some(len))
124     }
125
126     #[inline]
127     fn fold<Acc, Fold>(mut self, init: Acc, mut fold: Fold) -> Acc
128     where
129         Fold: FnMut(Acc, Self::Item) -> Acc,
130     {
131         let data = &mut self.data;
132         (&mut self.alive)
133             .try_fold::<_, _, Result<_, !>>(init, |acc, idx| {
134                 // SAFETY: idx is obtained by folding over the `alive` range, which implies the
135                 // value is currently considered alive but as the range is being consumed each value
136                 // we read here will only be read once and then considered dead.
137                 Ok(fold(acc, unsafe { data.get_unchecked(idx).assume_init_read() }))
138             })
139             .unwrap()
140     }
141
142     fn count(self) -> usize {
143         self.len()
144     }
145
146     fn last(mut self) -> Option<Self::Item> {
147         self.next_back()
148     }
149 }
150
151 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
152 impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
153     fn next_back(&mut self) -> Option<Self::Item> {
154         // Get the next index from the back.
155         //
156         // Decreasing `alive.end` by 1 maintains the invariant regarding
157         // `alive`. However, due to this change, for a short time, the alive
158         // zone is not `data[alive]` anymore, but `data[alive.start..=idx]`.
159         self.alive.next_back().map(|idx| {
160             // Read the element from the array.
161             // SAFETY: `idx` is an index into the former "alive" region of the
162             // array. Reading this element means that `data[idx]` is regarded as
163             // dead now (i.e. do not touch). As `idx` was the end of the
164             // alive-zone, the alive zone is now `data[alive]` again, restoring
165             // all invariants.
166             unsafe { self.data.get_unchecked(idx).assume_init_read() }
167         })
168     }
169 }
170
171 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
172 impl<T, const N: usize> Drop for IntoIter<T, N> {
173     fn drop(&mut self) {
174         // SAFETY: This is safe: `as_mut_slice` returns exactly the sub-slice
175         // of elements that have not been moved out yet and that remain
176         // to be dropped.
177         unsafe { ptr::drop_in_place(self.as_mut_slice()) }
178     }
179 }
180
181 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
182 impl<T, const N: usize> ExactSizeIterator for IntoIter<T, N> {
183     fn len(&self) -> usize {
184         // Will never underflow due to the invariant `alive.start <=
185         // alive.end`.
186         self.alive.end - self.alive.start
187     }
188     fn is_empty(&self) -> bool {
189         self.alive.is_empty()
190     }
191 }
192
193 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
194 impl<T, const N: usize> FusedIterator for IntoIter<T, N> {}
195
196 // The iterator indeed reports the correct length. The number of "alive"
197 // elements (that will still be yielded) is the length of the range `alive`.
198 // This range is decremented in length in either `next` or `next_back`. It is
199 // always decremented by 1 in those methods, but only if `Some(_)` is returned.
200 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
201 unsafe impl<T, const N: usize> TrustedLen for IntoIter<T, N> {}
202
203 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
204 impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
205     fn clone(&self) -> Self {
206         // Note, we don't really need to match the exact same alive range, so
207         // we can just clone into offset 0 regardless of where `self` is.
208         let mut new = Self { data: MaybeUninit::uninit_array(), alive: 0..0 };
209
210         // Clone all alive elements.
211         for (src, dst) in iter::zip(self.as_slice(), &mut new.data) {
212             // Write a clone into the new array, then update its alive range.
213             // If cloning panics, we'll correctly drop the previous items.
214             dst.write(src.clone());
215             new.alive.end += 1;
216         }
217
218         new
219     }
220 }
221
222 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
223 impl<T: fmt::Debug, const N: usize> fmt::Debug for IntoIter<T, N> {
224     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
225         // Only print the elements that were not yielded yet: we cannot
226         // access the yielded elements anymore.
227         f.debug_tuple("IntoIter").field(&self.as_slice()).finish()
228     }
229 }