Auto merge of #87168 - the8472:flatten-len, r=scottmcm

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.

resolves #87094

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,30 @@ impl<I, U, F> FusedIterator for FlatMap<I, U, F>
 {
 }
 
+#[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],
+{
+}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, T, I, F, const N: usize> TrustedLen for FlatMap<I, &'a [T; N], F>
+where
+    I: TrustedLen,
+    F: FnMut(I::Item) -> &'a [T; N],
+{
+}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, T, I, F, const N: usize> TrustedLen for FlatMap<I, &'a mut [T; N], F>
+where
+    I: TrustedLen,
+    F: FnMut(I::Item) -> &'a mut [T; N],
+{
+}
+
 /// An iterator that flattens one level of nesting in an iterator of things
 /// that can be turned into iterators.
 ///
@@ -230,6 +254,14 @@ impl<I, U> FusedIterator for Flatten<I>
 {
 }
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I> TrustedLen for Flatten<I>
+where
+    I: TrustedLen,
+    <I as Iterator>::Item: TrustedConstSize,
+{
+}
+
 /// Real logic of both `Flatten` and `FlatMap` which simply delegate to
 /// this type.
 #[derive(Clone, Debug)]
@@ -282,6 +314,17 @@ fn size_hint(&self) -> (usize, Option<usize>) {
         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 ConstSizeIntoIterator>::size() {
+            let (lower, upper) = self.iter.size_hint();
+
+            let lower = lower.saturating_mul(fixed_size).saturating_add(lo);
+            let upper =
+                try { fhi?.checked_add(bhi?)?.checked_add(fixed_size.checked_mul(upper?)?)? };
+
+            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 +487,52 @@ fn flatten<T: IntoIterator, Acc>(
         init
     }
 }
+
+trait ConstSizeIntoIterator: IntoIterator {
+    // FIXME(#31844): convert to an associated const once specialization supports that
+    fn size() -> Option<usize>;
+}
+
+impl<T> ConstSizeIntoIterator for T
+where
+    T: IntoIterator,
+{
+    #[inline]
+    default fn size() -> Option<usize> {
+        None
+    }
+}
+
+impl<T, const N: usize> ConstSizeIntoIterator for [T; N] {
+    #[inline]
+    fn size() -> Option<usize> {
+        Some(N)
+    }
+}
+
+impl<T, const N: usize> ConstSizeIntoIterator for &[T; N] {
+    #[inline]
+    fn size() -> Option<usize> {
+        Some(N)
+    }
+}
+
+impl<T, const N: usize> ConstSizeIntoIterator for &mut [T; N] {
+    #[inline]
+    fn size() -> Option<usize> {
+        Some(N)
+    }
+}
+
+#[doc(hidden)]
+#[unstable(feature = "std_internals", issue = "none")]
+// FIXME(#20400): Instead of this helper trait there should be multiple impl TrustedLen for Flatten<>
+//   blocks with different bounds on Iterator::Item but the compiler erroneously considers them overlapping
+pub unsafe trait TrustedConstSize: IntoIterator {}
+
+#[unstable(feature = "std_internals", issue = "none")]
+unsafe impl<T, const N: usize> TrustedConstSize for [T; N] {}
+#[unstable(feature = "std_internals", issue = "none")]
+unsafe impl<T, const N: usize> TrustedConstSize for &'_ [T; N] {}
+#[unstable(feature = "std_internals", issue = "none")]
+unsafe impl<T, const N: usize> TrustedConstSize for &'_ mut [T; N] {}

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

@@ -1,4 +1,5 @@
 use super::*;
+use core::array;
 use core::iter::*;
 
 #[test]
@@ -109,3 +110,42 @@ 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 mut a = [(); 10];
+    let mut b = [(); 10];
+
+    let iter = array::IntoIter::new([&mut a, &mut b]).flatten();
+    assert_trusted_len(&iter);
+    assert_eq!(iter.size_hint(), (20, Some(20)));
+    core::mem::drop(iter);
+
+    let iter = array::IntoIter::new([&a, &b]).flatten();
+    assert_trusted_len(&iter);
+    assert_eq!(iter.size_hint(), (20, Some(20)));
+
+    let iter = [(), (), ()].iter().flat_map(|_| [(); 1000]);
+    assert_trusted_len(&iter);
+    assert_eq!(iter.size_hint(), (3000, Some(3000)));
+
+    let iter = [(), ()].iter().flat_map(|_| &a);
+    assert_trusted_len(&iter);
+    assert_eq!(iter.size_hint(), (20, Some(20)));
+}