Skip to content

Simplify array::IntoIter #75271

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 8, 2020
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
144 changes: 55 additions & 89 deletions library/core/src/array/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,69 +56,55 @@ impl<T, const N: usize> IntoIter<T, N> {

// FIXME(LukasKalbertodt): actually use `mem::transmute` here, once it
// works with const generics:
// `mem::transmute::<[T; {N}], [MaybeUninit<T>; {N}]>(array)`
// `mem::transmute::<[T; N], [MaybeUninit<T>; N]>(array)`
//
// Until then, we do it manually here. We first create a bitwise copy
// but cast the pointer so that it is treated as a different type. Then
// we forget `array` so that it is not dropped.
let data = unsafe {
let data = ptr::read(&array as *const [T; N] as *const [MaybeUninit<T>; N]);
// Until then, we can use `mem::transmute_copy` to create a bitwise copy
// as a different type, then forget `array` so that it is not dropped.
unsafe {
let iter = Self { data: mem::transmute_copy(&array), alive: 0..N };
mem::forget(array);
data
};

Self { data, alive: 0..N }
iter
}
}

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

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

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> Iterator for IntoIter<T, N> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
if self.alive.start == self.alive.end {
return None;
}

// Bump start index.
// Get the next index from the front.
//
// From the check above we know that `alive.start != alive.end`.
// Combine this with the invariant `alive.start <= alive.end`, we know
// that `alive.start < alive.end`. Increasing `alive.start` by 1
// maintains the invariant regarding `alive`. However, due to this
// change, for a short time, the alive zone is not `data[alive]`
// anymore, but `data[idx..alive.end]`.
let idx = self.alive.start;
self.alive.start += 1;

// Read the element from the array.
// SAFETY: This is safe: `idx` is an index
// into the "alive" region of the array. Reading this element means
// that `data[idx]` is regarded as dead now (i.e. do not touch). As
// `idx` was the start of the alive-zone, the alive zone is now
// `data[alive]` again, restoring all invariants.
let out = unsafe { self.data.get_unchecked(idx).read() };

Some(out)
// Increasing `alive.start` by 1 maintains the invariant regarding
// `alive`. However, due to this change, for a short time, the alive
// zone is not `data[alive]` anymore, but `data[idx..alive.end]`.
self.alive.next().map(|idx| {
// Read the element from the array.
// SAFETY: `idx` is an index into the former "alive" region of the
// array. Reading this element means that `data[idx]` is regarded as
// dead now (i.e. do not touch). As `idx` was the start of the
// alive-zone, the alive zone is now `data[alive]` again, restoring
// all invariants.
unsafe { self.data.get_unchecked(idx).read() }
})
}

fn size_hint(&self) -> (usize, Option<usize>) {
Expand All @@ -138,33 +124,20 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
fn next_back(&mut self) -> Option<Self::Item> {
if self.alive.start == self.alive.end {
return None;
}

// Decrease end index.
// Get the next index from the back.
//
// From the check above we know that `alive.start != alive.end`.
// Combine this with the invariant `alive.start <= alive.end`, we know
// that `alive.start < alive.end`. As `alive.start` cannot be negative,
// `alive.end` is at least 1, meaning that we can safely decrement it
// by one. This also maintains the invariant `alive.start <=
// alive.end`. However, due to this change, for a short time, the alive
// zone is not `data[alive]` anymore, but `data[alive.start..alive.end
// + 1]`.
self.alive.end -= 1;

// Read the element from the array.
// SAFETY: This is safe: `alive.end` is an
// index into the "alive" region of the array. Compare the previous
// comment that states that the alive region is
// `data[alive.start..alive.end + 1]`. Reading this element means that
// `data[alive.end]` is regarded as dead now (i.e. do not touch). As
// `alive.end` was the end of the alive-zone, the alive zone is now
// `data[alive]` again, restoring all invariants.
let out = unsafe { self.data.get_unchecked(self.alive.end).read() };

Some(out)
// Decreasing `alive.end` by 1 maintains the invariant regarding
// `alive`. However, due to this change, for a short time, the alive
// zone is not `data[alive]` anymore, but `data[alive.start..=idx]`.
self.alive.next_back().map(|idx| {
// Read the element from the array.
// SAFETY: `idx` is an index into the former "alive" region of the
// array. Reading this element means that `data[idx]` is regarded as
// dead now (i.e. do not touch). As `idx` was the end of the
// alive-zone, the alive zone is now `data[alive]` again, restoring
// all invariants.
unsafe { self.data.get_unchecked(idx).read() }
})
}
}

Expand Down Expand Up @@ -203,26 +176,19 @@ unsafe impl<T, const N: usize> TrustedLen for IntoIter<T, N> {}
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
fn clone(&self) -> Self {
// SAFETY: each point of unsafety is documented inside the unsafe block
unsafe {
// This creates a new uninitialized array. Note that the `assume_init`
// refers to the array, not the individual elements. And it is Ok if
// the array is in an uninitialized state as all elements may be
// uninitialized (all bit patterns are valid). Compare the
// `MaybeUninit` docs for more information.
let mut new_data: [MaybeUninit<T>; N] = MaybeUninit::uninit().assume_init();

// Clone all alive elements.
for idx in self.alive.clone() {
// The element at `idx` in the old array is alive, so we can
// safely call `get_ref()`. We then clone it, and write the
// clone into the new array.
let clone = self.data.get_unchecked(idx).get_ref().clone();
new_data.get_unchecked_mut(idx).write(clone);
}

Self { data: new_data, alive: self.alive.clone() }
// Note, we don't really need to match the exact same alive range, so
// we can just clone into offset 0 regardless of where `self` is.
let mut new = Self { data: MaybeUninit::uninit_array(), alive: 0..0 };

// Clone all alive elements.
for (src, dst) in self.as_slice().iter().zip(&mut new.data) {
// Write a clone into the new array, then update its alive range.
// If cloning panics, we'll correctly drop the previous items.
dst.write(src.clone());
new.alive.end += 1;
}

new
}
}

Expand Down