diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index 3315d346596..48880a4d91a 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/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 FusedIterator for FlatMap { } +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for FlatMap +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 +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 +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 FusedIterator for Flatten { } +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for Flatten +where + I: TrustedLen, + ::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) { 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) = <::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( init } } + +trait ConstSizeIntoIterator: IntoIterator { + // FIXME(#31844): convert to an associated const once specialization supports that + fn size() -> Option; +} + +impl ConstSizeIntoIterator for T +where + T: IntoIterator, +{ + #[inline] + default fn size() -> Option { + None + } +} + +impl ConstSizeIntoIterator for [T; N] { + #[inline] + fn size() -> Option { + Some(N) + } +} + +impl ConstSizeIntoIterator for &[T; N] { + #[inline] + fn size() -> Option { + Some(N) + } +} + +impl ConstSizeIntoIterator for &mut [T; N] { + #[inline] + fn size() -> Option { + 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 TrustedConstSize for [T; N] {} +#[unstable(feature = "std_internals", issue = "none")] +unsafe impl TrustedConstSize for &'_ [T; N] {} +#[unstable(feature = "std_internals", issue = "none")] +unsafe impl TrustedConstSize for &'_ mut [T; N] {} diff --git a/library/core/tests/iter/adapters/flatten.rs b/library/core/tests/iter/adapters/flatten.rs index 4bbae6947bf..aaac39c2979 100644 --- a/library/core/tests/iter/adapters/flatten.rs +++ b/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) {} + 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))); +}