From bd1c39dc6c474971f6f0f3bc155661b574764c55 Mon Sep 17 00:00:00 2001 From: The8472 Date: Thu, 15 Jul 2021 22:59:26 +0200 Subject: [PATCH 1/4] 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. --- library/core/src/iter/adapters/flatten.rs | 49 ++++++++++++++++++++- library/core/tests/iter/adapters/flatten.rs | 24 ++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index 3315d34659611..0240d30363a57 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,14 @@ where { } +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for FlatMap +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 TrustedLen for Flatten where + I: Iterator + 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) = <::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; +} + +impl ConstSizeIterable for T { + #[inline] + default fn size() -> Option { + None + } +} + +impl ConstSizeIterable for [T; N] { + #[inline] + fn size() -> Option { + Some(N) + } +} diff --git a/library/core/tests/iter/adapters/flatten.rs b/library/core/tests/iter/adapters/flatten.rs index 4bbae6947bf66..9fa0ff4c92153 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,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) {} + 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))); +} From 18a034f97e8f69c8086301dad5ade9759883d775 Mon Sep 17 00:00:00 2001 From: The8472 Date: Fri, 16 Jul 2021 19:17:30 +0200 Subject: [PATCH 2/4] rename specializing trait to ConstSizeIntoIterator --- library/core/src/iter/adapters/flatten.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index 0240d30363a57..a6dc633a664ac 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/library/core/src/iter/adapters/flatten.rs @@ -297,7 +297,7 @@ where 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 ConstSizeIterable>::size() { + 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); @@ -474,18 +474,18 @@ where } } -trait ConstSizeIterable { +trait ConstSizeIntoIterator: IntoIterator { fn size() -> Option; } -impl ConstSizeIterable for T { +impl ConstSizeIntoIterator for T where T: IntoIterator { #[inline] default fn size() -> Option { None } } -impl ConstSizeIterable for [T; N] { +impl ConstSizeIntoIterator for [T; N] { #[inline] fn size() -> Option { Some(N) From 8dd903cc774e7376f4c66e7940fae8a420b25123 Mon Sep 17 00:00:00 2001 From: The8472 Date: Fri, 16 Jul 2021 20:30:53 +0200 Subject: [PATCH 3/4] 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. --- library/core/src/iter/adapters/flatten.rs | 55 +++++++++++++++++++-- library/core/tests/iter/adapters/flatten.rs | 16 ++++++ 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 a6dc633a664ac..9e0f6d2905d09 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 +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. /// @@ -239,8 +255,10 @@ where } #[unstable(feature = "trusted_len", issue = "37572")] -unsafe impl TrustedLen for Flatten where - I: Iterator + TrustedLen +unsafe impl TrustedLen for Flatten +where + I: TrustedLen, + ::Item: TrustedConstSize, { } @@ -475,10 +493,14 @@ where } trait ConstSizeIntoIterator: IntoIterator { + // FIXME(#31844): convert to an associated const once specialization supports that fn size() -> Option; } -impl ConstSizeIntoIterator for T where T: IntoIterator { +impl ConstSizeIntoIterator for T +where + T: IntoIterator, +{ #[inline] default fn size() -> Option { None @@ -491,3 +513,30 @@ impl ConstSizeIntoIterator for [T; N] { 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 9fa0ff4c92153..aaac39c297933 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))); } From c3ac8d8b8683fdb7f247e9107397ec8948ad4568 Mon Sep 17 00:00:00 2001 From: The8472 Date: Mon, 19 Jul 2021 20:21:54 +0200 Subject: [PATCH 4/4] replace Option combinators with try block --- library/core/src/iter/adapters/flatten.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index 9e0f6d2905d09..48880a4d91a57 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/library/core/src/iter/adapters/flatten.rs @@ -319,12 +319,8 @@ where 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(); + let upper = + try { fhi?.checked_add(bhi?)?.checked_add(fixed_size.checked_mul(upper?)?)? }; return (lower, upper); }