Skip to content

Commit bd1c39d

Browse files
committed
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.
1 parent 0a6c636 commit bd1c39d

File tree

2 files changed

+72
-1
lines changed

2 files changed

+72
-1
lines changed

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

+48-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use crate::fmt;
2-
use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map};
2+
use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map, TrustedLen};
33
use crate::ops::Try;
44

55
/// An iterator that maps each element to an iterator, and yields the elements
@@ -114,6 +114,14 @@ where
114114
{
115115
}
116116

117+
#[unstable(feature = "trusted_len", issue = "37572")]
118+
unsafe impl<T, I, F, const N: usize> TrustedLen for FlatMap<I, [T; N], F>
119+
where
120+
I: TrustedLen,
121+
F: FnMut(I::Item) -> [T; N],
122+
{
123+
}
124+
117125
/// An iterator that flattens one level of nesting in an iterator of things
118126
/// that can be turned into iterators.
119127
///
@@ -230,6 +238,12 @@ where
230238
{
231239
}
232240

241+
#[unstable(feature = "trusted_len", issue = "37572")]
242+
unsafe impl<T, I, const N: usize> TrustedLen for Flatten<I> where
243+
I: Iterator<Item = [T; N]> + TrustedLen
244+
{
245+
}
246+
233247
/// Real logic of both `Flatten` and `FlatMap` which simply delegate to
234248
/// this type.
235249
#[derive(Clone, Debug)]
@@ -282,6 +296,21 @@ where
282296
let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint);
283297
let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint);
284298
let lo = flo.saturating_add(blo);
299+
300+
if let Some(fixed_size) = <<I as Iterator>::Item as ConstSizeIterable>::size() {
301+
let (lower, upper) = self.iter.size_hint();
302+
303+
let lower = lower.saturating_mul(fixed_size).saturating_add(lo);
304+
let upper = upper.and_then(|i| i.checked_mul(fixed_size));
305+
let upper = fhi
306+
.zip_with(bhi, usize::checked_add)
307+
.flatten()
308+
.zip_with(upper, usize::checked_add)
309+
.flatten();
310+
311+
return (lower, upper);
312+
}
313+
285314
match (self.iter.size_hint(), fhi, bhi) {
286315
((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)),
287316
_ => (lo, None),
@@ -444,3 +473,21 @@ where
444473
init
445474
}
446475
}
476+
477+
trait ConstSizeIterable {
478+
fn size() -> Option<usize>;
479+
}
480+
481+
impl<T> ConstSizeIterable for T {
482+
#[inline]
483+
default fn size() -> Option<usize> {
484+
None
485+
}
486+
}
487+
488+
impl<T, const N: usize> ConstSizeIterable for [T; N] {
489+
#[inline]
490+
fn size() -> Option<usize> {
491+
Some(N)
492+
}
493+
}

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

+24
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use super::*;
2+
use core::array;
23
use core::iter::*;
34

45
#[test]
@@ -109,3 +110,26 @@ fn test_double_ended_flatten() {
109110
assert_eq!(it.next(), None);
110111
assert_eq!(it.next_back(), None);
111112
}
113+
114+
#[test]
115+
fn test_trusted_len_flatten() {
116+
fn assert_trusted_len<T: TrustedLen>(_: &T) {}
117+
let mut iter = array::IntoIter::new([[0; 3]; 4]).flatten();
118+
assert_trusted_len(&iter);
119+
120+
assert_eq!(iter.size_hint(), (12, Some(12)));
121+
iter.next();
122+
assert_eq!(iter.size_hint(), (11, Some(11)));
123+
iter.next_back();
124+
assert_eq!(iter.size_hint(), (10, Some(10)));
125+
126+
let iter = array::IntoIter::new([[(); usize::MAX]; 1]).flatten();
127+
assert_eq!(iter.size_hint(), (usize::MAX, Some(usize::MAX)));
128+
129+
let iter = array::IntoIter::new([[(); usize::MAX]; 2]).flatten();
130+
assert_eq!(iter.size_hint(), (usize::MAX, None));
131+
132+
let iter = [(), (), ()].iter().flat_map(|_| [(); 1000]);
133+
assert_trusted_len(&iter);
134+
assert_eq!(iter.size_hint(), (3000, Some(3000)));
135+
}

0 commit comments

Comments
 (0)