implement ConstSizeIntoIterator for &[T;N] in addition to [T;N]

Due to #20400 the corresponding TrustedLen impls need a helper trait
instead of directly adding `Item = &[T;N]` bounds.
Since TrustedLen is a public trait this in turn means
the helper trait needs to be public. Since it's just a workaround
for a compiler deficit it's marked hidden, unstable and unsafe.

2 files changed, 68 insertions, 3 deletions

diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs
index a6dc633a664..9e0f6d2905d 100644
--- a/library/core/src/iter/adapters/flatten.rs
+++ b/library/core/src/iter/adapters/flatten.rs
@@ -122,6 +122,22 @@ where
 {
 }
 
+#[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.
 ///
@@ -239,8 +255,10 @@ 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
+unsafe impl<I> TrustedLen for Flatten<I>
+where
+    I: TrustedLen,
+    <I as Iterator>::Item: TrustedConstSize,
 {
 }
 
@@ -475,10 +493,14 @@ where
 }
 
 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 {
+impl<T> ConstSizeIntoIterator for T
+where
+    T: IntoIterator,
+{
     #[inline]
     default fn size() -> Option<usize> {
         None
@@ -491,3 +513,30 @@ impl<T, const N: usize> ConstSizeIntoIterator for [T; N] {
         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] {}
diff --git a/library/core/tests/iter/adapters/flatten.rs b/library/core/tests/iter/adapters/flatten.rs
index 9fa0ff4c921..aaac39c2979 100644
--- a/library/core/tests/iter/adapters/flatten.rs
+++ b/library/core/tests/iter/adapters/flatten.rs
@@ -129,7 +129,23 @@ fn test_trusted_len_flatten() {
     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)));
 }