Skip to main content

arbitrary/foreign/core/
array.rs

1use {
2    crate::{size_hint, Arbitrary, Result, Unstructured},
3    core::{
4        array,
5        mem::{self, MaybeUninit},
6        ptr,
7    },
8};
9
10/// Helper to safely create arrays since the standard library doesn't
11/// provide one yet. Shouldn't be necessary in the future.
12struct ArrayGuard<T, const N: usize> {
13    dst: *mut T,
14    initialized: usize,
15}
16
17impl<T, const N: usize> Drop for ArrayGuard<T, N> {
18    fn drop(&mut self) {
19        debug_assert!(self.initialized <= N);
20        let initialized_part = ptr::slice_from_raw_parts_mut(self.dst, self.initialized);
21        // SAFETY:
22        // Contract from `ptr::drop_in_place`: the pointer must be valid for reads and
23        // writes, and must be properly aligned.
24        // Evidence:
25        // - `self.dst` is derived from `MaybeUninit<[T; N]>` on the stack, which remains
26        //   alive and allocated because the guard is a local variable dropped before
27        //   the stack frame is destroyed.
28        // - The alignment of `self.dst` matches that of `[T; N]`, which is aligned for `T`.
29        // - Only the first `self.initialized` elements are dropped, which have been
30        //   fully initialized by `ptr::write` in `try_create_array`.
31        // Therefore the contract of `drop_in_place` is discharged.
32        unsafe {
33            ptr::drop_in_place(initialized_part);
34        }
35    }
36}
37
38fn try_create_array<F, T, const N: usize>(mut cb: F) -> Result<[T; N]>
39where
40    F: FnMut(usize) -> Result<T>,
41{
42    let mut array: MaybeUninit<[T; N]> = MaybeUninit::uninit();
43    let array_ptr = array.as_mut_ptr();
44    let dst = array_ptr as *mut T;
45    let mut guard: ArrayGuard<T, N> = ArrayGuard {
46        dst,
47        initialized: 0,
48    };
49    for idx in 0..N {
50        // SAFETY: `dst` is a valid pointer to the start of the `[T; N]` array.
51        // `idx` is within `0..N`, so `dst.add(idx)` is within the bounds of the allocation.
52        // The pointer is properly aligned for `T`.
53        unsafe {
54            let value_ptr = dst.add(idx);
55            ptr::write(value_ptr, cb(idx)?);
56        }
57        guard.initialized += 1;
58    }
59    unsafe {
60        mem::forget(guard);
61        // SAFETY:
62        // Contract from `MaybeUninit::assume_init`: the value must be fully initialized.
63        // Evidence: the loop executes exactly `N` times, successfully writing a value
64        // to each of the `N` indices in the array. Since the loop completed without
65        // returning early or panicking, all elements of the array are initialized.
66        Ok(array.assume_init())
67    }
68}
69
70impl<'a, T, const N: usize> Arbitrary<'a> for [T; N]
71where
72    T: Arbitrary<'a>,
73{
74    #[inline]
75    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
76        try_create_array(|_| <T as Arbitrary<'a>>::arbitrary(u))
77    }
78
79    #[inline]
80    fn arbitrary_take_rest(mut u: Unstructured<'a>) -> Result<Self> {
81        let mut array = Self::arbitrary(&mut u)?;
82        if let Some(last) = array.last_mut() {
83            *last = Arbitrary::arbitrary_take_rest(u)?;
84        }
85        Ok(array)
86    }
87
88    #[inline]
89    fn size_hint(depth: usize) -> (usize, Option<usize>) {
90        Self::try_size_hint(depth).unwrap_or_default()
91    }
92
93    #[inline]
94    fn try_size_hint(depth: usize) -> Result<(usize, Option<usize>), crate::MaxRecursionReached> {
95        let hint = <T as Arbitrary>::try_size_hint(depth)?;
96        Ok(size_hint::and_all(&array::from_fn::<_, N, _>(|_| hint)))
97    }
98}