summaryrefslogtreecommitdiff
path: root/vendor/itertools/src/next_array.rs
diff options
context:
space:
mode:
Diffstat (limited to 'vendor/itertools/src/next_array.rs')
-rw-r--r--vendor/itertools/src/next_array.rs269
1 files changed, 269 insertions, 0 deletions
diff --git a/vendor/itertools/src/next_array.rs b/vendor/itertools/src/next_array.rs
new file mode 100644
index 00000000..86480b19
--- /dev/null
+++ b/vendor/itertools/src/next_array.rs
@@ -0,0 +1,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);
+ }
+ }
+}
generated by cgit v1.2.3 (git 2.39.1) at 2026-02-02 22:11:27 +0000