Skip to content

Commit

Permalink
Add repeat and resize(_with) for arrays
Browse files Browse the repository at this point in the history
```rust
mod array {
    pub fn repeat<T: Clone, const N: usize>(value: T) -> [T; N];
}

impl [T; N] {
    pub fn resize<const NEW_LEN: usize>(self, value: T) -> [T; NEW_LEN] where T: Clone;
    pub fn resize_with<const NEW_LEN: usize>(self, mut f: impl FnMut() -> T) -> [T; NEW_LEN];
}
```

I tried to do `resize_with` as just
```rust
collect_into_array(IntoIter::new(self).chain(iter::repeat_with(f)))
```
but unfortunately that optimized very poorly.

As a result there's a bunch of refactoring in here to move the `Guard` struct that was previously private to `try_collect_into_array` over to its own module.  It's still very much internal -- it's visible only to the `core::array` module, with no intention of making this the anointed public thing.
  • Loading branch information
scottmcm committed Dec 6, 2021
1 parent 0fb1c37 commit 5a784e6
Show file tree
Hide file tree
Showing 4 changed files with 395 additions and 25 deletions.
96 changes: 96 additions & 0 deletions library/core/src/array/guard.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
use crate::mem::{forget, replace, MaybeUninit};
use crate::ptr;

/// The internal-use drop guard for implementing array methods.
///
/// This is free to be changed whenever. Its purpose is not to provide a
/// beautiful safe interface, but to make the unsafe details of `super`'s
/// other methods slightly more obvious and have reduced code duplication.
pub struct Guard<'a, T, const N: usize> {
array_mut: &'a mut [MaybeUninit<T>; N],
initialized: usize,
}

impl<'a, T, const N: usize> Guard<'a, T, N> {
#[inline]
pub fn new(buffer: &'a mut [MaybeUninit<T>; N]) -> Self {
Self { array_mut: buffer, initialized: 0 }
}

#[inline]
pub fn len(&self) -> usize {
self.initialized
}

/// Initialize the next item
///
/// # Safety
///
/// Requires `self.len() < N`.
#[inline]
pub unsafe fn push_unchecked(&mut self, value: T) {
debug_assert!(self.len() < N);
// SAFETY: The precondition means we have space
unsafe {
self.array_mut.get_unchecked_mut(self.initialized).write(value);
}
self.initialized += 1;
}

/// Initialize the next `CHUNK` item(s)
///
/// # Safety
///
/// Requires `self.len() + CHUNK <= N`.
#[inline]
pub unsafe fn push_chunk_unchecked<const CHUNK: usize>(&mut self, chunk: [T; CHUNK]) {
assert!(CHUNK <= N);
debug_assert!(N - self.len() >= CHUNK);
// SAFETY: The precondition means we have space
unsafe {
// Since we're going to write multiple items, make sure not to do so
// via a `&mut MaybeUninit<T>`, as that would violate stacked borrows.
let first = self.array_mut.as_mut_ptr();
let p = first.add(self.initialized).cast();
ptr::write(p, chunk);
}
self.initialized += CHUNK;
}

/// Read the whole buffer as an initialized array.
///
/// This always de-initializes the original buffer -- even if `T: Copy`.
///
/// # Safety
///
/// Requires `self.len() == N`.
#[inline]
pub unsafe fn into_array_unchecked(self) -> [T; N] {
debug_assert_eq!(self.len(), N);

// This tells LLVM and MIRI that we don't care about the buffer after,
// and the extra `undef` write is trivial for it to optimize away.
let buffer = replace(self.array_mut, MaybeUninit::uninit_array());

// SAFETY: the condition above asserts that all elements are
// initialized.
let out = unsafe { MaybeUninit::array_assume_init(buffer) };

forget(self);

out
}
}

impl<T, const N: usize> Drop for Guard<'_, T, N> {
fn drop(&mut self) {
debug_assert!(self.initialized <= N);

// SAFETY: this slice will contain only initialized objects.
unsafe {
crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(
&mut self.array_mut.get_unchecked_mut(..self.initialized),
));
}
}
}
21 changes: 21 additions & 0 deletions library/core/src/array/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,27 @@ impl<T, const N: usize> IntoIter<T, N> {
MaybeUninit::slice_assume_init_mut(slice)
}
}

/// Returns the remaining `M` items of this iterator as an array.
///
/// If there are more than `M` items left in this iterator, the rest of
/// them will be leaked. So that should be avoided, but it's sound.
///
/// # Safety
///
/// There must be at least `M` items remaining.
pub(crate) unsafe fn into_array_unchecked<const M: usize>(self) -> [T; M] {
debug_assert!(self.len() >= M);

// SAFETY: The precondition of at least M items left means that
// there are enough valid contiguous items for the `read`.
let array: [T; M] = unsafe { ptr::read(self.as_slice().as_ptr().cast()) };

// We better not run any drops for the items we're about to return.
mem::forget(self);

array
}
}

#[stable(feature = "array_value_iter_impls", since = "1.40.0")]
Expand Down
Loading

0 comments on commit 5a784e6

Please sign in to comment.