1 files changed, 148 insertions, 0 deletions
diff --git a/vendor/unarray/src/from_iter.rs b/vendor/unarray/src/from_iter.rs
new file mode 100644
index 000000000..e96275bb2
--- /dev/null
+++ b/vendor/unarray/src/from_iter.rs
@@ -0,0 +1,148 @@
+use crate::{mark_initialized, uninit_buf};
+use core::iter::FromIterator;
+
+/// A wrapper type to collect an [`Iterator`] into an array
+///
+/// ```
+/// # use unarray::*;
+/// let iter = vec![1, 2, 3].into_iter();
+/// let ArrayFromIter(array) = iter.collect();
+///
+/// assert_eq!(array, Some([1, 2, 3]));
+/// ```
+/// Since iterators don't carry compile-time information about their length (even
+/// [`core::iter::ExactSizeIterator`] only provides this at runtime), collecting may fail if the
+/// iterator doesn't yield **exactly** `N` elements:
+/// ```
+/// use unarray::*;
+/// let too_many = vec![1, 2, 3, 4].into_iter();
+/// let ArrayFromIter::<i32, 3>(option) = too_many.collect();
+/// assert!(option.is_none());
+///
+/// let too_few = vec![1, 2].into_iter();
+/// let ArrayFromIter::<i32, 3>(option) = too_few.collect();
+/// assert!(option.is_none());
+/// ```
+pub struct ArrayFromIter<T, const N: usize>(pub Option<[T; N]>);
+
+impl<T, const N: usize> FromIterator<T> for ArrayFromIter<T, N> {
+    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
+        let mut buffer = uninit_buf::<T, N>();
+        let mut iter = iter.into_iter();
+        let mut buf_iter = buffer.iter_mut();
+        let mut num_written = 0;
+
+        loop {
+            let item = iter.next();
+            let slot = buf_iter.next();
+
+            match (item, slot) {
+                (Some(item), Some(slot)) => {
+                    num_written += 1;
+                    slot.write(item);
+                }
+                // error case, we should free the previous ones and continue
+                (Some(_), None) | (None, Some(_)) => {
+                    // SAFETY:
+                    // There are `num_written` elements fully initialized, so we can safely drop
+                    // them
+                    buffer
+                        .iter_mut()
+                        .take(num_written)
+                        .for_each(|slot| unsafe { slot.assume_init_drop() });
+
+                    return Self(None);
+                }
+                // SAFETY
+                // If this is reached, every prior iteration of the loop has matched
+                // (Some(_), Some(_)). As such, both iterators have yielded the same number of
+                // elements, so every slot has been written to
+                (None, None) => return Self(Some(unsafe { mark_initialized(buffer) })),
+            };
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use core::{
+        convert::TryInto,
+        sync::atomic::{AtomicUsize, Ordering},
+    };
+
+    use crate::testing::vec_strategy;
+    use proptest::{prop_assert, prop_assert_eq};
+    use test_strategy::proptest;
+
+    use super::*;
+
+    #[test]
+    fn can_collect_array_from_iter() {
+        let iter = vec![1, 2, 3].into_iter();
+
+        let ArrayFromIter(array) = iter.collect();
+        assert_eq!(array.unwrap(), [1, 2, 3]);
+    }
+
+    #[test]
+    fn fails_if_incorrect_number_of_elements() {
+        let iter = [1, 2, 3].iter();
+        let ArrayFromIter::<_, 4>(array) = iter.collect();
+        assert!(array.is_none());
+
+        let iter = [1, 2, 3].iter();
+        let ArrayFromIter::<_, 2>(array) = iter.collect();
+        assert!(array.is_none());
+    }
+
+    const LEN: usize = 100;
+    const SHORT_LEN: usize = LEN - 1;
+    const LONG_LEN: usize = LEN + 1;
+
+    #[derive(Clone)]
+    struct IncrementOnDrop<'a>(&'a AtomicUsize);
+    impl Drop for IncrementOnDrop<'_> {
+        fn drop(&mut self) {
+            self.0.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    #[test]
+    fn doesnt_leak_too_long() {
+        let drop_count = 0.into();
+        let ArrayFromIter::<_, 3>(_) = vec![IncrementOnDrop(&drop_count); 4].into_iter().collect();
+        // since it failed, all 4 should be dropped
+        assert_eq!(drop_count.load(Ordering::Relaxed), 4);
+    }
+
+    #[test]
+    fn doesnt_leak_too_short() {
+        let drop_count = 0.into();
+        let ArrayFromIter::<_, 3>(_) = vec![IncrementOnDrop(&drop_count); 2].into_iter().collect();
+        // since it failed, both should be dropped
+        assert_eq!(drop_count.load(Ordering::Relaxed), 2);
+    }
+
+    #[proptest]
+    #[cfg_attr(miri, ignore)]
+    fn undersized_proptest(#[strategy(vec_strategy(LEN))] vec: Vec<String>) {
+        let ArrayFromIter::<String, SHORT_LEN>(array) = vec.into_iter().collect();
+        prop_assert!(array.is_none());
+    }
+
+    #[proptest]
+    #[cfg_attr(miri, ignore)]
+    fn oversized_proptest(#[strategy(vec_strategy(LEN))] vec: Vec<String>) {
+        let ArrayFromIter::<String, LONG_LEN>(array) = vec.into_iter().collect();
+        prop_assert!(array.is_none());
+    }
+
+    #[proptest]
+    #[cfg_attr(miri, ignore)]
+    fn just_right_proptest(#[strategy(vec_strategy(LEN))] vec: Vec<String>) {
+        let expected: [String; LEN] = vec.clone().try_into().unwrap();
+        let ArrayFromIter(array) = vec.into_iter().collect();
+        prop_assert_eq!(array.unwrap(), expected);
+    }
+}