implement TrustedLen for Flatten/FlatMap if the U: IntoIterator == [T; N]

This only works if arrays are passed directly instead of array iterators
because we need to be sure that they have not been advanced before
Flatten does its size calculation.

Author: the8472

library/core/src/iter/adapters/flatten.rs

@@ -1,5 +1,5 @@
 use crate::fmt;
-use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map};
+use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map, TrustedLen};
 use crate::ops::Try;
 
 /// An iterator that maps each element to an iterator, and yields the elements
@@ -114,6 +114,14 @@ where
 {
 }
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<T, I, F, const N: usize> TrustedLen for FlatMap<I, [T; N], F>
+where
+    I: TrustedLen,
+    F: FnMut(I::Item) -> [T; N],
+{
+}
+
 /// An iterator that flattens one level of nesting in an iterator of things
 /// that can be turned into iterators.
 ///
@@ -230,6 +238,12 @@ where
 {
 }
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<T, I, const N: usize> TrustedLen for Flatten<I> where
+    I: Iterator<Item = [T; N]> + TrustedLen
+{
+}
+
 /// Real logic of both `Flatten` and `FlatMap` which simply delegate to
 /// this type.
 #[derive(Clone, Debug)]
@@ -282,6 +296,21 @@ where
         let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint);
         let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint);
         let lo = flo.saturating_add(blo);
+
+        if let Some(fixed_size) = <<I as Iterator>::Item as ConstSizeIterable>::size() {
+            let (lower, upper) = self.iter.size_hint();
+
+            let lower = lower.saturating_mul(fixed_size).saturating_add(lo);
+            let upper = upper.and_then(|i| i.checked_mul(fixed_size));
+            let upper = fhi
+                .zip_with(bhi, usize::checked_add)
+                .flatten()
+                .zip_with(upper, usize::checked_add)
+                .flatten();
+
+            return (lower, upper);
+        }
+
         match (self.iter.size_hint(), fhi, bhi) {
             ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)),
             _ => (lo, None),
@@ -444,3 +473,21 @@ where
         init
     }
 }
+
+trait ConstSizeIterable {
+    fn size() -> Option<usize>;
+}
+
+impl<T> ConstSizeIterable for T {
+    #[inline]
+    default fn size() -> Option<usize> {
+        None
+    }
+}
+
+impl<T, const N: usize> ConstSizeIterable for [T; N] {
+    #[inline]
+    fn size() -> Option<usize> {
+        Some(N)
+    }
+}

library/core/tests/iter/adapters/flatten.rs

@@ -1,4 +1,5 @@
 use super::*;
+use core::array;
 use core::iter::*;
 
 #[test]
@@ -109,3 +110,26 @@ fn test_double_ended_flatten() {
     assert_eq!(it.next(), None);
     assert_eq!(it.next_back(), None);
 }
+
+#[test]
+fn test_trusted_len_flatten() {
+    fn assert_trusted_len<T: TrustedLen>(_: &T) {}
+    let mut iter = array::IntoIter::new([[0; 3]; 4]).flatten();
+    assert_trusted_len(&iter);
+
+    assert_eq!(iter.size_hint(), (12, Some(12)));
+    iter.next();
+    assert_eq!(iter.size_hint(), (11, Some(11)));
+    iter.next_back();
+    assert_eq!(iter.size_hint(), (10, Some(10)));
+
+    let iter = array::IntoIter::new([[(); usize::MAX]; 1]).flatten();
+    assert_eq!(iter.size_hint(), (usize::MAX, Some(usize::MAX)));
+
+    let iter = array::IntoIter::new([[(); usize::MAX]; 2]).flatten();
+    assert_eq!(iter.size_hint(), (usize::MAX, None));
+
+    let iter = [(), (), ()].iter().flat_map(|_| [(); 1000]);
+    assert_trusted_len(&iter);
+    assert_eq!(iter.size_hint(), (3000, Some(3000)));
+}