Skip to content

Commit ceedf1d

Browse files
borsbors
committed
Auto merge of #75271 - cuviper:array-iter, r=LukasKalbertodt
Simplify array::IntoIter - Initialization can use `transmute_copy` to do the bitwise copy. - `as_slice` can use `get_unchecked` and `MaybeUninit::slice_get_ref`, and `as_mut_slice` can do similar. - `next` and `next_back` can use the corresponding `Range` methods. - `Clone` doesn't need any unsafety, and we can dynamically update the new range to get partial drops if `T::clone` panics. r? @LukasKalbertodt
2 parents 1facd4a + a2cfc74 commit ceedf1d

File tree

1 file changed

+55
-89
lines changed

1 file changed

+55
-89
lines changed

‎library/core/src/array/iter.rs‎

Lines changed: 55 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -56,69 +56,55 @@ impl<T, const N: usize> IntoIter<T, N> {
5656

5757
// FIXME(LukasKalbertodt): actually use `mem::transmute` here, once it
5858
// works with const generics:
59-
// `mem::transmute::<[T; {N}], [MaybeUninit<T>; {N}]>(array)`
59+
// `mem::transmute::<[T; N], [MaybeUninit<T>; N]>(array)`
6060
//
61-
// Until then, we do it manually here. We first create a bitwise copy
62-
// but cast the pointer so that it is treated as a different type. Then
63-
// we forget `array` so that it is not dropped.
64-
let data = unsafe {
65-
let data = ptr::read(&array as *const [T; N] as *const [MaybeUninit<T>; N]);
61+
// Until then, we can use `mem::transmute_copy` to create a bitwise copy
62+
// as a different type, then forget `array` so that it is not dropped.
63+
unsafe {
64+
let iter = Self { data: mem::transmute_copy(&array), alive: 0..N };
6665
mem::forget(array);
67-
data
68-
};
69-
70-
Self { data, alive: 0..N }
66+
iter
67+
}
7168
}
7269

7370
/// Returns an immutable slice of all elements that have not been yielded
7471
/// yet.
7572
fn as_slice(&self) -> &[T] {
76-
let slice = &self.data[self.alive.clone()];
77-
// SAFETY: This transmute is safe. As mentioned in `new`, `MaybeUninit` retains
78-
// the size and alignment of `T`. Furthermore, we know that all
79-
// elements within `alive` are properly initialized.
80-
unsafe { mem::transmute::<&[MaybeUninit<T>], &[T]>(slice) }
73+
// SAFETY: We know that all elements within `alive` are properly initialized.
74+
unsafe {
75+
let slice = self.data.get_unchecked(self.alive.clone());
76+
MaybeUninit::slice_get_ref(slice)
77+
}
8178
}
8279

8380
/// Returns a mutable slice of all elements that have not been yielded yet.
8481
fn as_mut_slice(&mut self) -> &mut [T] {
85-
// This transmute is safe, same as in `as_slice` above.
86-
let slice = &mut self.data[self.alive.clone()];
87-
// SAFETY: This transmute is safe. As mentioned in `new`, `MaybeUninit` retains
88-
// the size and alignment of `T`. Furthermore, we know that all
89-
// elements within `alive` are properly initialized.
90-
unsafe { mem::transmute::<&mut [MaybeUninit<T>], &mut [T]>(slice) }
82+
// SAFETY: We know that all elements within `alive` are properly initialized.
83+
unsafe {
84+
let slice = self.data.get_unchecked_mut(self.alive.clone());
85+
MaybeUninit::slice_get_mut(slice)
86+
}
9187
}
9288
}
9389

9490
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
9591
impl<T, const N: usize> Iterator for IntoIter<T, N> {
9692
type Item = T;
9793
fn next(&mut self) -> Option<Self::Item> {
98-
if self.alive.start == self.alive.end {
99-
return None;
100-
}
101-
102-
// Bump start index.
94+
// Get the next index from the front.
10395
//
104-
// From the check above we know that `alive.start != alive.end`.
105-
// Combine this with the invariant `alive.start <= alive.end`, we know
106-
// that `alive.start < alive.end`. Increasing `alive.start` by 1
107-
// maintains the invariant regarding `alive`. However, due to this
108-
// change, for a short time, the alive zone is not `data[alive]`
109-
// anymore, but `data[idx..alive.end]`.
110-
let idx = self.alive.start;
111-
self.alive.start += 1;
112-
113-
// Read the element from the array.
114-
// SAFETY: This is safe: `idx` is an index
115-
// into the "alive" region of the array. Reading this element means
116-
// that `data[idx]` is regarded as dead now (i.e. do not touch). As
117-
// `idx` was the start of the alive-zone, the alive zone is now
118-
// `data[alive]` again, restoring all invariants.
119-
let out = unsafe { self.data.get_unchecked(idx).read() };
120-
121-
Some(out)
96+
// Increasing `alive.start` by 1 maintains the invariant regarding
97+
// `alive`. However, due to this change, for a short time, the alive
98+
// zone is not `data[alive]` anymore, but `data[idx..alive.end]`.
99+
self.alive.next().map(|idx| {
100+
// Read the element from the array.
101+
// SAFETY: `idx` is an index into the former "alive" region of the
102+
// array. Reading this element means that `data[idx]` is regarded as
103+
// dead now (i.e. do not touch). As `idx` was the start of the
104+
// alive-zone, the alive zone is now `data[alive]` again, restoring
105+
// all invariants.
106+
unsafe { self.data.get_unchecked(idx).read() }
107+
})
122108
}
123109

124110
fn size_hint(&self) -> (usize, Option<usize>) {
@@ -138,33 +124,20 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
138124
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
139125
impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
140126
fn next_back(&mut self) -> Option<Self::Item> {
141-
if self.alive.start == self.alive.end {
142-
return None;
143-
}
144-
145-
// Decrease end index.
127+
// Get the next index from the back.
146128
//
147-
// From the check above we know that `alive.start != alive.end`.
148-
// Combine this with the invariant `alive.start <= alive.end`, we know
149-
// that `alive.start < alive.end`. As `alive.start` cannot be negative,
150-
// `alive.end` is at least 1, meaning that we can safely decrement it
151-
// by one. This also maintains the invariant `alive.start <=
152-
// alive.end`. However, due to this change, for a short time, the alive
153-
// zone is not `data[alive]` anymore, but `data[alive.start..alive.end
154-
// + 1]`.
155-
self.alive.end -= 1;
156-
157-
// Read the element from the array.
158-
// SAFETY: This is safe: `alive.end` is an
159-
// index into the "alive" region of the array. Compare the previous
160-
// comment that states that the alive region is
161-
// `data[alive.start..alive.end + 1]`. Reading this element means that
162-
// `data[alive.end]` is regarded as dead now (i.e. do not touch). As
163-
// `alive.end` was the end of the alive-zone, the alive zone is now
164-
// `data[alive]` again, restoring all invariants.
165-
let out = unsafe { self.data.get_unchecked(self.alive.end).read() };
166-
167-
Some(out)
129+
// Decreasing `alive.end` by 1 maintains the invariant regarding
130+
// `alive`. However, due to this change, for a short time, the alive
131+
// zone is not `data[alive]` anymore, but `data[alive.start..=idx]`.
132+
self.alive.next_back().map(|idx| {
133+
// Read the element from the array.
134+
// SAFETY: `idx` is an index into the former "alive" region of the
135+
// array. Reading this element means that `data[idx]` is regarded as
136+
// dead now (i.e. do not touch). As `idx` was the end of the
137+
// alive-zone, the alive zone is now `data[alive]` again, restoring
138+
// all invariants.
139+
unsafe { self.data.get_unchecked(idx).read() }
140+
})
168141
}
169142
}
170143

@@ -203,26 +176,19 @@ unsafe impl<T, const N: usize> TrustedLen for IntoIter<T, N> {}
203176
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
204177
impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
205178
fn clone(&self) -> Self {
206-
// SAFETY: each point of unsafety is documented inside the unsafe block
207-
unsafe {
208-
// This creates a new uninitialized array. Note that the `assume_init`
209-
// refers to the array, not the individual elements. And it is Ok if
210-
// the array is in an uninitialized state as all elements may be
211-
// uninitialized (all bit patterns are valid). Compare the
212-
// `MaybeUninit` docs for more information.
213-
let mut new_data: [MaybeUninit<T>; N] = MaybeUninit::uninit().assume_init();
214-
215-
// Clone all alive elements.
216-
for idx in self.alive.clone() {
217-
// The element at `idx` in the old array is alive, so we can
218-
// safely call `get_ref()`. We then clone it, and write the
219-
// clone into the new array.
220-
let clone = self.data.get_unchecked(idx).get_ref().clone();
221-
new_data.get_unchecked_mut(idx).write(clone);
222-
}
223-
224-
Self { data: new_data, alive: self.alive.clone() }
179+
// Note, we don't really need to match the exact same alive range, so
180+
// we can just clone into offset 0 regardless of where `self` is.
181+
let mut new = Self { data: MaybeUninit::uninit_array(), alive: 0..0 };
182+
183+
// Clone all alive elements.
184+
for (src, dst) in self.as_slice().iter().zip(&mut new.data) {
185+
// Write a clone into the new array, then update its alive range.
186+
// If cloning panics, we'll correctly drop the previous items.
187+
dst.write(src.clone());
188+
new.alive.end += 1;
225189
}
190+
191+
new
226192
}
227193
}
228194

0 commit comments

Comments
 (0)