From dc12cde9238874a7b4efb02c76fddbb54af50354 Mon Sep 17 00:00:00 2001
From: Caio <c410.f3r@gmail.com>
Date: Fri, 12 Feb 2021 08:42:27 -0300
Subject: [PATCH] Add const-generics feature

---
 src/lib.rs | 135 +++++++++++++++++++++++++++--------------------------
 1 file changed, 68 insertions(+), 67 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index f92850b..27985df 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -713,89 +713,90 @@ 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<T: Arbitrary> Arbitrary for [T; $n] {
-            fn arbitrary(u: &mut Unstructured<'_>) -> Result<[T; $n]> {
-                Ok([
-                    Arbitrary::arbitrary(u)?,
-                    $(<$ts as Arbitrary>::arbitrary(u)?),*
-                ])
-            }
-
-            #[allow(unused_mut)]
-            fn arbitrary_take_rest(mut u: Unstructured<'_>) -> Result<[T; $n]> {
-                $(let $as = $ts::arbitrary(&mut u)?;)*
-                let last = Arbitrary::arbitrary_take_rest(u)?;
-
-                Ok([
-                    $($as,)* last
-                ])
-            }
+struct ArrayGuard<T, const N: usize> {
+    dst: *mut T,
+    initialized: usize,
+}
 
-            #[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) ),*
-                ])
-            }
+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);
+        }
+    }
+}
 
-            #[allow(unused_mut)] // For the `[T; 1]` case.
-            fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
-                let mut i = 0;
-                let mut shrinkers = [
-                    self[i].shrink(),
-                    $({
-                        i += 1;
-                        let t: &$ts = &self[i];
-                        t.shrink()
-                    }),*
-                ];
-                Box::new(iter::from_fn(move || {
-                    let mut i = 0;
-                    Some([
-                        shrinkers[i].next()?,
-                        $({
-                            i += 1;
-                            let t: $ts = shrinkers[i].next()?;
-                            t
-                        }),*
-                    ])
-                }))
-            }
+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<T: Arbitrary> Arbitrary for [T; 0] {
-    fn arbitrary(_: &mut Unstructured<'_>) -> Result<[T; 0]> {
-        Ok([])
+impl<T, const N: usize> Arbitrary for [T; N]
+where
+    T: Arbitrary,
+{
+    fn arbitrary(u: &mut Unstructured<'_>) -> Result<[T; N]> {
+        try_create_array(|_| <T as Arbitrary>::arbitrary(u))
     }
 
-    fn arbitrary_take_rest(_: Unstructured<'_>) -> Result<[T; 0]> {
-        Ok([])
+    fn arbitrary_take_rest(mut u: Unstructured<'_>) -> Result<[T; N]> {
+        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)
+        }))
     }
 
     fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
-        Box::new(iter::from_fn(|| None))
+        let mut shrinkers =
+            create_array::<_, Box<dyn Iterator<Item = T>>, N>(|idx| self[idx].shrink());
+        Box::new(iter::from_fn(move || {
+            try_create_array(|idx| shrinkers[idx].next().ok_or(Error::NotEnoughData)).ok()
+        }))
     }
 }
 
-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) }
-
 fn shrink_collection<'a, T, A: Arbitrary>(
     entries: impl Iterator<Item = T>,
     f: impl Fn(&T) -> Box<dyn Iterator<Item = A>>,