Skip to content

Commit

Permalink
Fix bug in iter::Chain::size_hint
Browse files Browse the repository at this point in the history
  • Loading branch information
timvermeulen committed Aug 18, 2019
1 parent ea52be4 commit ec54340
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 ec54340

Please sign in to comment.