From 84bc38af793f37f0d73a76f0e393a5ccf519c9ae Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Thu, 3 Dec 2020 13:57:18 +0300
Subject: [PATCH] Add Iterator::collect_array method

With const-generics, we can now collect exactly `N` elements out of
the iterator.

Very similar API existed in the itertools crate for quite some time:

https://docs.rs/itertools/0.9.0/itertools/trait.Itertools.html#method.collect_tuple
---
 library/core/src/iter/traits/iterator.rs | 73 ++++++++++++++++++++++++
 1 file changed, 73 insertions(+)

diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs
index 19484bfd0419f..6d60af18bacef 100644
--- a/library/core/src/iter/traits/iterator.rs
+++ b/library/core/src/iter/traits/iterator.rs
@@ -3,7 +3,9 @@
 // can't split that into multiple files.
 
 use crate::cmp::{self, Ordering};
+use crate::mem::{self, MaybeUninit};
 use crate::ops::{Add, ControlFlow, Try};
+use crate::ptr;
 
 use super::super::TrustedRandomAccess;
 use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
@@ -1670,6 +1672,77 @@ pub trait Iterator {
         FromIterator::from_iter(self)
     }
 
+    /// Collects all items from the iterator into an array of a specific size.
+    ///
+    /// If the number of elements inside the iterator is exactly equal to the
+    /// array size, then the array is returned inside `Some`, otherwise `None`
+    /// is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(iter_collect_array)]
+    ///
+    /// let xs = [1, 2, 3];
+    /// let iter = xs.iter().copied();
+    /// let [_a, b, _c] = iter.clone().collect_array().unwrap();
+    /// assert_eq!(b, 2);
+    ///
+    /// match iter.collect_array() {
+    ///     Some([_, _]) => panic!("Didn't expect to see two only elements"),
+    ///     None => (),
+    /// }
+    /// ```
+    #[inline]
+    #[unstable(reason = "new API", issue = "none", feature = "iter_collect_array")]
+    fn collect_array<const N: usize>(self) -> Option<[Self::Item; N]>
+    where
+        Self: Sized,
+    {
+        struct Buf<T, const N: usize> {
+            // Safety invariant: first `len` items are initialized.
+            items: [MaybeUninit<T>; N],
+            len: usize,
+        }
+        impl<T, const N: usize> Buf<T, N> {
+            fn new() -> Buf<T, N> {
+                Buf { items: MaybeUninit::uninit_array(), len: 0 }
+            }
+            fn push(&mut self, item: T) -> Option<()> {
+                let slot = self.items.get_mut(self.len)?;
+                slot.write(item);
+                self.len += 1;
+                Some(())
+            }
+            fn into_array(mut self) -> Option<[T; N]> {
+                if self.len != N {
+                    return None;
+                }
+                self.len = 0;
+                let res =
+                    // SAFETY: `len` field invariant + the guard above.
+                    unsafe { mem::transmute_copy::<[MaybeUninit<T>; N], [T; N]>(&self.items) };
+                Some(res)
+            }
+        }
+
+        impl<T, const N: usize> Drop for Buf<T, N> {
+            fn drop(&mut self) {
+                // SAFETY: `len` field invariant.
+                unsafe {
+                    let slice = MaybeUninit::slice_assume_init_mut(&mut self.items[..self.len]);
+                    ptr::drop_in_place(slice);
+                }
+            }
+        }
+
+        let mut buf: Buf<Self::Item, { N }> = Buf::new();
+        for elem in self {
+            buf.push(elem)?;
+        }
+        buf.into_array()
+    }
+
     /// Consumes an iterator, creating two collections from it.
     ///
     /// The predicate passed to `partition()` can return `true`, or `false`.