itertools/next_array.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269
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);
}
}
}