270 lines
9.1 KiB
Rust
270 lines
9.1 KiB
Rust
use core::mem::{self, MaybeUninit};
|
|
|
|
/// An array of at most `N` elements.
|
|
struct ArrayBuilder<T, const N: usize> {
|
|
/// The (possibly uninitialized) elements of the `ArrayBuilder`.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// The elements of `arr[..len]` are valid `T`s.
|
|
arr: [MaybeUninit<T>; N],
|
|
|
|
/// The number of leading elements of `arr` that are valid `T`s, len <= N.
|
|
len: usize,
|
|
}
|
|
|
|
impl<T, const N: usize> ArrayBuilder<T, N> {
|
|
/// Initializes a new, empty `ArrayBuilder`.
|
|
pub fn new() -> Self {
|
|
// SAFETY: The safety invariant of `arr` trivially holds for `len = 0`.
|
|
Self {
|
|
arr: [(); N].map(|_| MaybeUninit::uninit()),
|
|
len: 0,
|
|
}
|
|
}
|
|
|
|
/// Pushes `value` onto the end of the array.
|
|
///
|
|
/// # Panics
|
|
///
|
|
/// This panics if `self.len >= N`.
|
|
#[inline(always)]
|
|
pub fn push(&mut self, value: T) {
|
|
// PANICS: This will panic if `self.len >= N`.
|
|
let place = &mut self.arr[self.len];
|
|
// SAFETY: The safety invariant of `self.arr` applies to elements at
|
|
// indices `0..self.len` — not to the element at `self.len`. Writing to
|
|
// the element at index `self.len` therefore does not violate the safety
|
|
// invariant of `self.arr`. Even if this line panics, we have not
|
|
// created any intermediate invalid state.
|
|
*place = MaybeUninit::new(value);
|
|
// Lemma: `self.len < N`. By invariant, `self.len <= N`. Above, we index
|
|
// into `self.arr`, which has size `N`, at index `self.len`. If `self.len == N`
|
|
// at that point, that index would be out-of-bounds, and the index
|
|
// operation would panic. Thus, `self.len != N`, and since `self.len <= N`,
|
|
// that means that `self.len < N`.
|
|
//
|
|
// PANICS: Since `self.len < N`, and since `N <= usize::MAX`,
|
|
// `self.len + 1 <= usize::MAX`, and so `self.len += 1` will not
|
|
// overflow. Overflow is the only panic condition of `+=`.
|
|
//
|
|
// SAFETY:
|
|
// - We are required to uphold the invariant that `self.len <= N`.
|
|
// Since, by the preceding lemma, `self.len < N` at this point in the
|
|
// code, `self.len += 1` results in `self.len <= N`.
|
|
// - We are required to uphold the invariant that `self.arr[..self.len]`
|
|
// are valid instances of `T`. Since this invariant already held when
|
|
// this method was called, and since we only increment `self.len`
|
|
// by 1 here, we only need to prove that the element at
|
|
// `self.arr[self.len]` (using the value of `self.len` before incrementing)
|
|
// is valid. Above, we construct `place` to point to `self.arr[self.len]`,
|
|
// and then initialize `*place` to `MaybeUninit::new(value)`, which is
|
|
// a valid `T` by construction.
|
|
self.len += 1;
|
|
}
|
|
|
|
/// Consumes the elements in the `ArrayBuilder` and returns them as an array
|
|
/// `[T; N]`.
|
|
///
|
|
/// If `self.len() < N`, this returns `None`.
|
|
pub fn take(&mut self) -> Option<[T; N]> {
|
|
if self.len == N {
|
|
// SAFETY: Decreasing the value of `self.len` cannot violate the
|
|
// safety invariant on `self.arr`.
|
|
self.len = 0;
|
|
|
|
// SAFETY: Since `self.len` is 0, `self.arr` may safely contain
|
|
// uninitialized elements.
|
|
let arr = mem::replace(&mut self.arr, [(); N].map(|_| MaybeUninit::uninit()));
|
|
|
|
Some(arr.map(|v| {
|
|
// SAFETY: We know that all elements of `arr` are valid because
|
|
// we checked that `len == N`.
|
|
unsafe { v.assume_init() }
|
|
}))
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<T, const N: usize> AsMut<[T]> for ArrayBuilder<T, N> {
|
|
fn as_mut(&mut self) -> &mut [T] {
|
|
let valid = &mut self.arr[..self.len];
|
|
// SAFETY: By invariant on `self.arr`, the elements of `self.arr` at
|
|
// indices `0..self.len` are in a valid state. Since `valid` references
|
|
// only these elements, the safety precondition of
|
|
// `slice_assume_init_mut` is satisfied.
|
|
unsafe { slice_assume_init_mut(valid) }
|
|
}
|
|
}
|
|
|
|
impl<T, const N: usize> Drop for ArrayBuilder<T, N> {
|
|
// We provide a non-trivial `Drop` impl, because the trivial impl would be a
|
|
// no-op; `MaybeUninit<T>` has no innate awareness of its own validity, and
|
|
// so it can only forget its contents. By leveraging the safety invariant of
|
|
// `self.arr`, we do know which elements of `self.arr` are valid, and can
|
|
// selectively run their destructors.
|
|
fn drop(&mut self) {
|
|
// SAFETY:
|
|
// - by invariant on `&mut [T]`, `self.as_mut()` is:
|
|
// - valid for reads and writes
|
|
// - properly aligned
|
|
// - non-null
|
|
// - the dropped `T` are valid for dropping; they do not have any
|
|
// additional library invariants that we've violated
|
|
// - no other pointers to `valid` exist (since we're in the context of
|
|
// `drop`)
|
|
unsafe { core::ptr::drop_in_place(self.as_mut()) }
|
|
}
|
|
}
|
|
|
|
/// Assuming all the elements are initialized, get a mutable slice to them.
|
|
///
|
|
/// # Safety
|
|
///
|
|
/// The caller guarantees that the elements `T` referenced by `slice` are in a
|
|
/// valid state.
|
|
unsafe fn slice_assume_init_mut<T>(slice: &mut [MaybeUninit<T>]) -> &mut [T] {
|
|
// SAFETY: Casting `&mut [MaybeUninit<T>]` to `&mut [T]` is sound, because
|
|
// `MaybeUninit<T>` is guaranteed to have the same size, alignment and ABI
|
|
// as `T`, and because the caller has guaranteed that `slice` is in the
|
|
// valid state.
|
|
unsafe { &mut *(slice as *mut [MaybeUninit<T>] as *mut [T]) }
|
|
}
|
|
|
|
/// Equivalent to `it.next_array()`.
|
|
pub(crate) fn next_array<I, const N: usize>(it: &mut I) -> Option<[I::Item; N]>
|
|
where
|
|
I: Iterator,
|
|
{
|
|
let mut builder = ArrayBuilder::new();
|
|
for _ in 0..N {
|
|
builder.push(it.next()?);
|
|
}
|
|
builder.take()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use super::ArrayBuilder;
|
|
|
|
#[test]
|
|
fn zero_len_take() {
|
|
let mut builder = ArrayBuilder::<(), 0>::new();
|
|
let taken = builder.take();
|
|
assert_eq!(taken, Some([(); 0]));
|
|
}
|
|
|
|
#[test]
|
|
#[should_panic]
|
|
fn zero_len_push() {
|
|
let mut builder = ArrayBuilder::<(), 0>::new();
|
|
builder.push(());
|
|
}
|
|
|
|
#[test]
|
|
fn push_4() {
|
|
let mut builder = ArrayBuilder::<(), 4>::new();
|
|
assert_eq!(builder.take(), None);
|
|
|
|
builder.push(());
|
|
assert_eq!(builder.take(), None);
|
|
|
|
builder.push(());
|
|
assert_eq!(builder.take(), None);
|
|
|
|
builder.push(());
|
|
assert_eq!(builder.take(), None);
|
|
|
|
builder.push(());
|
|
assert_eq!(builder.take(), Some([(); 4]));
|
|
}
|
|
|
|
#[test]
|
|
fn tracked_drop() {
|
|
use std::panic::{catch_unwind, AssertUnwindSafe};
|
|
use std::sync::atomic::{AtomicU16, Ordering};
|
|
|
|
static DROPPED: AtomicU16 = AtomicU16::new(0);
|
|
|
|
#[derive(Debug, PartialEq)]
|
|
struct TrackedDrop;
|
|
|
|
impl Drop for TrackedDrop {
|
|
fn drop(&mut self) {
|
|
DROPPED.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
}
|
|
|
|
{
|
|
let builder = ArrayBuilder::<TrackedDrop, 0>::new();
|
|
assert_eq!(DROPPED.load(Ordering::Relaxed), 0);
|
|
drop(builder);
|
|
assert_eq!(DROPPED.load(Ordering::Relaxed), 0);
|
|
}
|
|
|
|
{
|
|
let mut builder = ArrayBuilder::<TrackedDrop, 2>::new();
|
|
builder.push(TrackedDrop);
|
|
assert_eq!(builder.take(), None);
|
|
assert_eq!(DROPPED.load(Ordering::Relaxed), 0);
|
|
drop(builder);
|
|
assert_eq!(DROPPED.swap(0, Ordering::Relaxed), 1);
|
|
}
|
|
|
|
{
|
|
let mut builder = ArrayBuilder::<TrackedDrop, 2>::new();
|
|
builder.push(TrackedDrop);
|
|
builder.push(TrackedDrop);
|
|
assert!(matches!(builder.take(), Some(_)));
|
|
assert_eq!(DROPPED.swap(0, Ordering::Relaxed), 2);
|
|
drop(builder);
|
|
assert_eq!(DROPPED.load(Ordering::Relaxed), 0);
|
|
}
|
|
|
|
{
|
|
let mut builder = ArrayBuilder::<TrackedDrop, 2>::new();
|
|
|
|
builder.push(TrackedDrop);
|
|
builder.push(TrackedDrop);
|
|
|
|
assert!(catch_unwind(AssertUnwindSafe(|| {
|
|
builder.push(TrackedDrop);
|
|
}))
|
|
.is_err());
|
|
|
|
assert_eq!(DROPPED.load(Ordering::Relaxed), 1);
|
|
|
|
drop(builder);
|
|
|
|
assert_eq!(DROPPED.swap(0, Ordering::Relaxed), 3);
|
|
}
|
|
|
|
{
|
|
let mut builder = ArrayBuilder::<TrackedDrop, 2>::new();
|
|
|
|
builder.push(TrackedDrop);
|
|
builder.push(TrackedDrop);
|
|
|
|
assert!(catch_unwind(AssertUnwindSafe(|| {
|
|
builder.push(TrackedDrop);
|
|
}))
|
|
.is_err());
|
|
|
|
assert_eq!(DROPPED.load(Ordering::Relaxed), 1);
|
|
|
|
assert!(matches!(builder.take(), Some(_)));
|
|
|
|
assert_eq!(DROPPED.load(Ordering::Relaxed), 3);
|
|
|
|
builder.push(TrackedDrop);
|
|
builder.push(TrackedDrop);
|
|
|
|
assert!(matches!(builder.take(), Some(_)));
|
|
|
|
assert_eq!(DROPPED.swap(0, Ordering::Relaxed), 5);
|
|
}
|
|
}
|
|
}
|