Skip to content

Commit 5a784e6

Browse files
committed
Add repeat and resize(_with) for arrays
```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.
1 parent 0fb1c37 commit 5a784e6

File tree

4 files changed

+395
-25
lines changed

4 files changed

+395
-25
lines changed

library/core/src/array/guard.rs

+96
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
use crate::mem::{forget, replace, MaybeUninit};
2+
use crate::ptr;
3+
4+
/// The internal-use drop guard for implementing array methods.
5+
///
6+
/// This is free to be changed whenever. Its purpose is not to provide a
7+
/// beautiful safe interface, but to make the unsafe details of `super`'s
8+
/// other methods slightly more obvious and have reduced code duplication.
9+
pub struct Guard<'a, T, const N: usize> {
10+
array_mut: &'a mut [MaybeUninit<T>; N],
11+
initialized: usize,
12+
}
13+
14+
impl<'a, T, const N: usize> Guard<'a, T, N> {
15+
#[inline]
16+
pub fn new(buffer: &'a mut [MaybeUninit<T>; N]) -> Self {
17+
Self { array_mut: buffer, initialized: 0 }
18+
}
19+
20+
#[inline]
21+
pub fn len(&self) -> usize {
22+
self.initialized
23+
}
24+
25+
/// Initialize the next item
26+
///
27+
/// # Safety
28+
///
29+
/// Requires `self.len() < N`.
30+
#[inline]
31+
pub unsafe fn push_unchecked(&mut self, value: T) {
32+
debug_assert!(self.len() < N);
33+
// SAFETY: The precondition means we have space
34+
unsafe {
35+
self.array_mut.get_unchecked_mut(self.initialized).write(value);
36+
}
37+
self.initialized += 1;
38+
}
39+
40+
/// Initialize the next `CHUNK` item(s)
41+
///
42+
/// # Safety
43+
///
44+
/// Requires `self.len() + CHUNK <= N`.
45+
#[inline]
46+
pub unsafe fn push_chunk_unchecked<const CHUNK: usize>(&mut self, chunk: [T; CHUNK]) {
47+
assert!(CHUNK <= N);
48+
debug_assert!(N - self.len() >= CHUNK);
49+
// SAFETY: The precondition means we have space
50+
unsafe {
51+
// Since we're going to write multiple items, make sure not to do so
52+
// via a `&mut MaybeUninit<T>`, as that would violate stacked borrows.
53+
let first = self.array_mut.as_mut_ptr();
54+
let p = first.add(self.initialized).cast();
55+
ptr::write(p, chunk);
56+
}
57+
self.initialized += CHUNK;
58+
}
59+
60+
/// Read the whole buffer as an initialized array.
61+
///
62+
/// This always de-initializes the original buffer -- even if `T: Copy`.
63+
///
64+
/// # Safety
65+
///
66+
/// Requires `self.len() == N`.
67+
#[inline]
68+
pub unsafe fn into_array_unchecked(self) -> [T; N] {
69+
debug_assert_eq!(self.len(), N);
70+
71+
// This tells LLVM and MIRI that we don't care about the buffer after,
72+
// and the extra `undef` write is trivial for it to optimize away.
73+
let buffer = replace(self.array_mut, MaybeUninit::uninit_array());
74+
75+
// SAFETY: the condition above asserts that all elements are
76+
// initialized.
77+
let out = unsafe { MaybeUninit::array_assume_init(buffer) };
78+
79+
forget(self);
80+
81+
out
82+
}
83+
}
84+
85+
impl<T, const N: usize> Drop for Guard<'_, T, N> {
86+
fn drop(&mut self) {
87+
debug_assert!(self.initialized <= N);
88+
89+
// SAFETY: this slice will contain only initialized objects.
90+
unsafe {
91+
crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(
92+
&mut self.array_mut.get_unchecked_mut(..self.initialized),
93+
));
94+
}
95+
}
96+
}

library/core/src/array/iter.rs

+21
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,27 @@ impl<T, const N: usize> IntoIter<T, N> {
104104
MaybeUninit::slice_assume_init_mut(slice)
105105
}
106106
}
107+
108+
/// Returns the remaining `M` items of this iterator as an array.
109+
///
110+
/// If there are more than `M` items left in this iterator, the rest of
111+
/// them will be leaked. So that should be avoided, but it's sound.
112+
///
113+
/// # Safety
114+
///
115+
/// There must be at least `M` items remaining.
116+
pub(crate) unsafe fn into_array_unchecked<const M: usize>(self) -> [T; M] {
117+
debug_assert!(self.len() >= M);
118+
119+
// SAFETY: The precondition of at least M items left means that
120+
// there are enough valid contiguous items for the `read`.
121+
let array: [T; M] = unsafe { ptr::read(self.as_slice().as_ptr().cast()) };
122+
123+
// We better not run any drops for the items we're about to return.
124+
mem::forget(self);
125+
126+
array
127+
}
107128
}
108129

109130
#[stable(feature = "array_value_iter_impls", since = "1.40.0")]

0 commit comments

Comments
 (0)