From 23c6602a170f99f9519c1c55b551afad5dd1366e Mon Sep 17 00:00:00 2001
From: Caio <c410.f3r@gmail.com>
Date: Thu, 25 Mar 2021 14:38:10 -0300
Subject: [PATCH 1/2] Add const-generics feature

---
 src/lib.rs | 106 ++++++++++++++++++++++++++++++++---------------------
 1 file changed, 65 insertions(+), 41 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 213614d..bb604ae 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -571,61 +571,85 @@ macro_rules! arbitrary_tuple {
 }
 arbitrary_tuple!(A B C D E F G H I J K L M N O P Q R S T U V W X Y Z);
 
-macro_rules! arbitrary_array {
-    {$n:expr, ($t:ident, $a:ident) $(($ts:ident, $as:ident))*} => {
-        arbitrary_array!{($n - 1), $(($ts, $as))*}
-
-        impl<'a, T: Arbitrary<'a>> Arbitrary<'a> for [T; $n] {
-            fn arbitrary(u: &mut Unstructured<'a>) -> Result<[T; $n]> {
-                Ok([
-                    Arbitrary::arbitrary(u)?,
-                    $(<$ts as Arbitrary>::arbitrary(u)?),*
-                ])
-            }
-
-            #[allow(unused_mut)]
-            fn arbitrary_take_rest(mut u: Unstructured<'a>) -> Result<[T; $n]> {
-                $(let $as = $ts::arbitrary(&mut u)?;)*
-                let last = Arbitrary::arbitrary_take_rest(u)?;
+struct ArrayGuard<T, const N: usize> {
+    dst: *mut T,
+    initialized: usize,
+}
 
-                Ok([
-                    $($as,)* last
-                ])
-            }
+impl<T, const N: usize> Drop for ArrayGuard<T, N> {
+    fn drop(&mut self) {
+        debug_assert!(self.initialized <= N);
+        let initialized_part = core::ptr::slice_from_raw_parts_mut(self.dst, self.initialized);
+        unsafe {
+            core::ptr::drop_in_place(initialized_part);
+        }
+    }
+}
 
-            #[inline]
-            fn size_hint(depth: usize) -> (usize, Option<usize>) {
-                crate::size_hint::and_all(&[
-                    <$t as Arbitrary>::size_hint(depth),
-                    $( <$ts as Arbitrary>::size_hint(depth) ),*
-                ])
-            }
+fn create_array<F, T, const N: usize>(mut cb: F) -> [T; N]
+where
+    F: FnMut(usize) -> T,
+{
+    let mut array: mem::MaybeUninit<[T; N]> = mem::MaybeUninit::uninit();
+    let mut guard: ArrayGuard<T, N> = ArrayGuard {
+        dst: array.as_mut_ptr() as _,
+        initialized: 0,
+    };
+    unsafe {
+        for (idx, value_ptr) in (&mut *array.as_mut_ptr()).iter_mut().enumerate() {
+            core::ptr::write(value_ptr, cb(idx));
+            guard.initialized += 1;
         }
+        mem::forget(guard);
+        array.assume_init()
+    }
+}
+
+fn try_create_array<F, T, const N: usize>(mut cb: F) -> Result<[T; N]>
+where
+    F: FnMut(usize) -> Result<T>,
+{
+    let mut array: mem::MaybeUninit<[T; N]> = mem::MaybeUninit::uninit();
+    let mut guard: ArrayGuard<T, N> = ArrayGuard {
+        dst: array.as_mut_ptr() as _,
+        initialized: 0,
     };
-    ($n: expr,) => {};
+    unsafe {
+        for (idx, value_ptr) in (&mut *array.as_mut_ptr()).iter_mut().enumerate() {
+            core::ptr::write(value_ptr, cb(idx)?);
+            guard.initialized += 1;
+        }
+        mem::forget(guard);
+        Ok(array.assume_init())
+    }
 }
 
-impl<'a, T: Arbitrary<'a>> Arbitrary<'a> for [T; 0] {
-    fn arbitrary(_: &mut Unstructured<'a>) -> Result<[T; 0]> {
-        Ok([])
+impl<'a, T, const N: usize> Arbitrary<'a> for [T; N]
+where
+    T: Arbitrary<'a>,
+{
+    #[inline]
+    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
+        try_create_array(|_| <T as Arbitrary<'a>>::arbitrary(u))
     }
 
-    fn arbitrary_take_rest(_: Unstructured<'a>) -> Result<[T; 0]> {
-        Ok([])
+    #[inline]
+    fn arbitrary_take_rest(mut u: Unstructured<'a>) -> Result<Self> {
+        let mut array = Self::arbitrary(&mut u)?;
+        if let Some(last) = array.last_mut() {
+            *last = Arbitrary::arbitrary_take_rest(u)?;
+        }
+        Ok(array)
     }
 
     #[inline]
-    fn size_hint(_: usize) -> (usize, Option<usize>) {
-        crate::size_hint::and_all(&[])
+    fn size_hint(d: usize) -> (usize, Option<usize>) {
+        crate::size_hint::and_all(&create_array::<_, (usize, Option<usize>), N>(|_| {
+            <T as Arbitrary>::size_hint(d)
+        }))
     }
 }
 
-arbitrary_array! { 32, (T, a) (T, b) (T, c) (T, d) (T, e) (T, f) (T, g) (T, h)
-(T, i) (T, j) (T, k) (T, l) (T, m) (T, n) (T, o) (T, p)
-(T, q) (T, r) (T, s) (T, u) (T, v) (T, w) (T, x) (T, y)
-(T, z) (T, aa) (T, ab) (T, ac) (T, ad) (T, ae) (T, af)
-(T, ag) }
-
 impl<'a> Arbitrary<'a> for &'a [u8] {
     fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
         let len = u.arbitrary_len::<u8>()?;

From f69268c30eea73df34ed968f9c66e85f3b8862f5 Mon Sep 17 00:00:00 2001
From: Caio <c410.f3r@gmail.com>
Date: Mon, 29 Mar 2021 18:14:32 -0300
Subject: [PATCH 2/2] Address comments

---
 src/lib.rs | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index bb604ae..b8458b4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -571,6 +571,8 @@ macro_rules! arbitrary_tuple {
 }
 arbitrary_tuple!(A B C D E F G H I J K L M N O P Q R S T U V W X Y Z);
 
+// Helper to safely create arrays since the standard library doesn't
+// provide one yet. Shouldn't be necessary in the future.
 struct ArrayGuard<T, const N: usize> {
     dst: *mut T,
     initialized: usize,
@@ -591,8 +593,10 @@ where
     F: FnMut(usize) -> T,
 {
     let mut array: mem::MaybeUninit<[T; N]> = mem::MaybeUninit::uninit();
+    let array_ptr = array.as_mut_ptr();
+    let dst = array_ptr as _;
     let mut guard: ArrayGuard<T, N> = ArrayGuard {
-        dst: array.as_mut_ptr() as _,
+        dst,
         initialized: 0,
     };
     unsafe {
@@ -610,8 +614,10 @@ where
     F: FnMut(usize) -> Result<T>,
 {
     let mut array: mem::MaybeUninit<[T; N]> = mem::MaybeUninit::uninit();
+    let array_ptr = array.as_mut_ptr();
+    let dst = array_ptr as _;
     let mut guard: ArrayGuard<T, N> = ArrayGuard {
-        dst: array.as_mut_ptr() as _,
+        dst,
         initialized: 0,
     };
     unsafe {