Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Produce better size_hints from Flatten & FlatMap #48544

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/libcore/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,8 @@ macro_rules! array_impls {
fn into_iter(self) -> Iter<'a, T> {
self.iter()
}

const SIZE_HINT: (usize, Option<usize>) = ($N, Some($N));
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand All @@ -208,6 +210,8 @@ macro_rules! array_impls {
fn into_iter(self) -> IterMut<'a, T> {
self.iter_mut()
}

const SIZE_HINT: (usize, Option<usize>) = ($N, Some($N));
}

// NOTE: some less important impls are omitted to reduce code bloat
Expand Down
20 changes: 15 additions & 5 deletions src/libcore/iter/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2611,11 +2611,21 @@ impl<I, U> Iterator for FlattenCompat<I, U>
fn size_hint(&self) -> (usize, Option<usize>) {
let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), |it| it.size_hint());
let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), |it| it.size_hint());
let lo = flo.saturating_add(blo);
match (self.iter.size_hint(), fhi, bhi) {
((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)),
_ => (lo, None)
}

let (ilo, ihi) = self.iter.size_hint();
let (clo, chi) = <I::Item as IntoIterator>::SIZE_HINT;
let mlo = ilo.saturating_mul(clo);
let mhi = match (ihi, chi) {
(Some(0), _) => Some(0),
(Some(a), Some(b)) => a.checked_mul(b),
_ => None,
};

let lo = flo.saturating_add(blo).saturating_add(mlo);
let combine_hi = || {
fhi?.checked_add(bhi?)?.checked_add(mhi?)
};
(lo, combine_hi())
}

#[inline]
Expand Down
22 changes: 22 additions & 0 deletions src/libcore/iter/traits.rs
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,28 @@ pub trait IntoIterator {
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
fn into_iter(self) -> Self::IntoIter;

/// A hint for the size of the iterators produced by this type.
///
/// This allows things like `.flatten()` to return a meaningful
/// `.size_hint()` before the actual iterators have been produced.
///
/// This is more useful than having it on `Iterator`, since finite
/// iterators always have a lower-bound size of `0` by definition.
/// A `&[T; N]`, in contrast, knows that its iterator will produce
/// exactly `N` items even though its iterator type is a normal
/// slice iterator that has a wide range of `.size_hint()`s.
///
/// # Examples
///
/// ```
/// #![feature(typelevel_size_hint)]
///
/// assert_eq!(<&'static [i32; 10] as IntoIterator>::SIZE_HINT, (10, Some(10)));
/// assert_eq!(<Option<String> as IntoIterator>::SIZE_HINT, (0, Some(1)));
/// ```
#[unstable(feature = "typelevel_size_hint", issue = "7777777")]
const SIZE_HINT: (usize, Option<usize>) = (0, None);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A style nit, but why are you declaring the const SIZE_HINTs beneath all of the method definitions? They kind of feel like they should be before method definitions?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had it up there originally, but then it felt like it was pretending to be too important, and that the things one actually needs to implement should be first. I'm happy to use either.

}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down
6 changes: 6 additions & 0 deletions src/libcore/option.rs
Original file line number Diff line number Diff line change
Expand Up @@ -953,6 +953,8 @@ impl<T> IntoIterator for Option<T> {
fn into_iter(self) -> IntoIter<T> {
IntoIter { inner: Item { opt: self } }
}

const SIZE_HINT: (usize, Option<usize>) = (0, Some(1));
}

#[stable(since = "1.4.0", feature = "option_iter")]
Expand All @@ -963,6 +965,8 @@ impl<'a, T> IntoIterator for &'a Option<T> {
fn into_iter(self) -> Iter<'a, T> {
self.iter()
}

const SIZE_HINT: (usize, Option<usize>) = (0, Some(1));
}

#[stable(since = "1.4.0", feature = "option_iter")]
Expand All @@ -973,6 +977,8 @@ impl<'a, T> IntoIterator for &'a mut Option<T> {
fn into_iter(self) -> IterMut<'a, T> {
self.iter_mut()
}

const SIZE_HINT: (usize, Option<usize>) = (0, Some(1));
}

#[stable(since = "1.12.0", feature = "option_from")]
Expand Down
6 changes: 6 additions & 0 deletions src/libcore/result.rs
Original file line number Diff line number Diff line change
Expand Up @@ -977,6 +977,8 @@ impl<T, E> IntoIterator for Result<T, E> {
fn into_iter(self) -> IntoIter<T> {
IntoIter { inner: self.ok() }
}

const SIZE_HINT: (usize, Option<usize>) = (0, Some(1));
}

#[stable(since = "1.4.0", feature = "result_iter")]
Expand All @@ -987,6 +989,8 @@ impl<'a, T, E> IntoIterator for &'a Result<T, E> {
fn into_iter(self) -> Iter<'a, T> {
self.iter()
}

const SIZE_HINT: (usize, Option<usize>) = (0, Some(1));
}

#[stable(since = "1.4.0", feature = "result_iter")]
Expand All @@ -997,6 +1001,8 @@ impl<'a, T, E> IntoIterator for &'a mut Result<T, E> {
fn into_iter(self) -> IterMut<'a, T> {
self.iter_mut()
}

const SIZE_HINT: (usize, Option<usize>) = (0, Some(1));
}

/////////////////////////////////////////////////////////////////////////////
Expand Down
22 changes: 22 additions & 0 deletions src/libcore/tests/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -885,6 +885,23 @@ fn test_iterator_flatten() {
i += 1;
}
assert_eq!(i, ys.len());

let aa = [ [1, 2, 3], [4, 5, 6], [7, 8, 9] ];
let mut i = aa.iter().flatten();
assert_eq!(i.size_hint(), (9, Some(9)));
// Ensure the frontiter and backiter are included
assert_eq!(i.next(), Some(&1));
assert_eq!(i.size_hint(), (8, Some(8)));
assert_eq!(i.next_back(), Some(&9));
assert_eq!(i.size_hint(), (7, Some(7)));

let vv = vec![ vec![1] ];
let mut i = vv.iter().flatten();
// Vec precludes a good static estimate
assert_eq!(i.size_hint(), (0, None));
i.next();
// Until it's empty, where we no longer need the estimate
assert_eq!(i.size_hint(), (0, Some(0)));
}

/// Test `Flatten::fold` with items already picked off the front and back,
Expand Down Expand Up @@ -1073,7 +1090,12 @@ fn test_iterator_size_hint() {
assert_eq!(vi.clone().scan(0, |_,_| Some(0)).size_hint(), (0, Some(10)));
assert_eq!(vi.clone().filter(|_| false).size_hint(), (0, Some(10)));
assert_eq!(vi.clone().map(|&i| i+1).size_hint(), (10, Some(10)));
assert_eq!(vi.clone().filter_map(|_| Some(1)).size_hint(), (0, Some(10)));
assert_eq!(vi.clone().flat_map(|_| Some(1)).size_hint(), (0, Some(10)));
assert_eq!(vi.filter_map(|_| Some(0)).size_hint(), (0, Some(10)));

let vv = [ [1, 2, 3], [4, 5, 6], [7, 8, 9] ];
assert_eq!(vv.iter().flatten().size_hint(), (9, Some(9)));
}

#[test]
Expand Down