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

Author: bors

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 @@ 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],
+{
+}
+
+#[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 @@ where
 {
 }
 
+#[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 @@ 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 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 @@ where
         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)));
+}