Skip to content

Commit

Permalink
Rollup merge of rust-lang#63691 - timvermeulen:chain-size-hint, r=sco…
Browse files Browse the repository at this point in the history
…ttmcm

Fix bug in iter::Chain::size_hint

`Chain::size_hint` currently ignores `self.state`, which means that the size hints of the underlying iterators are always combined regardless of the iteration state. This, of course, should only happen when the state is `ChainState::Both`.
  • Loading branch information
Centril committed Aug 20, 2019
2 parents 4dd3ce8 + ec54340 commit 9dfb7ae
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 8 deletions.
22 changes: 14 additions & 8 deletions src/libcore/iter/adapters/chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -173,17 +173,23 @@ impl<A, B> Iterator for Chain<A, B> where

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let (a_lower, a_upper) = self.a.size_hint();
let (b_lower, b_upper) = self.b.size_hint();
match self.state {
ChainState::Both => {
let (a_lower, a_upper) = self.a.size_hint();
let (b_lower, b_upper) = self.b.size_hint();

let lower = a_lower.saturating_add(b_lower);
let lower = a_lower.saturating_add(b_lower);

let upper = match (a_upper, b_upper) {
(Some(x), Some(y)) => x.checked_add(y),
_ => None
};
let upper = match (a_upper, b_upper) {
(Some(x), Some(y)) => x.checked_add(y),
_ => None
};

(lower, upper)
(lower, upper)
}
ChainState::Front => self.a.size_hint(),
ChainState::Back => self.b.size_hint(),
}
}
}

Expand Down
48 changes: 48 additions & 0 deletions src/libcore/tests/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,54 @@ fn test_iterator_chain_find() {
assert_eq!(iter.next(), None);
}

#[test]
fn test_iterator_chain_size_hint() {
struct Iter {
is_empty: bool,
}

impl Iterator for Iter {
type Item = ();

// alternates between `None` and `Some(())`
fn next(&mut self) -> Option<Self::Item> {
if self.is_empty {
self.is_empty = false;
None
} else {
self.is_empty = true;
Some(())
}
}

fn size_hint(&self) -> (usize, Option<usize>) {
if self.is_empty {
(0, Some(0))
} else {
(1, Some(1))
}
}
}

impl DoubleEndedIterator for Iter {
fn next_back(&mut self) -> Option<Self::Item> {
self.next()
}
}

// this chains an iterator of length 0 with an iterator of length 1,
// so after calling `.next()` once, the iterator is empty and the
// state is `ChainState::Back`. `.size_hint()` should now disregard
// the size hint of the left iterator
let mut iter = Iter { is_empty: true }.chain(once(()));
assert_eq!(iter.next(), Some(()));
assert_eq!(iter.size_hint(), (0, Some(0)));

let mut iter = once(()).chain(Iter { is_empty: true });
assert_eq!(iter.next_back(), Some(()));
assert_eq!(iter.size_hint(), (0, Some(0)));
}

#[test]
fn test_zip_nth() {
let xs = [0, 1, 2, 4, 5];
Expand Down

0 comments on commit 9dfb7ae

Please sign in to comment.