Skip to content

Commit

Permalink
Add ExactSizeIterator impl for iter::Chain
Browse files Browse the repository at this point in the history
  • Loading branch information
LukasKalbertodt committed Dec 12, 2019
1 parent a0d40f8 commit 722bba0
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 0 deletions.
30 changes: 30 additions & 0 deletions src/libcore/iter/adapters/chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,36 @@ impl<A, B> FusedIterator for Chain<A, B>
B: FusedIterator<Item=A::Item>,
{}

#[stable(feature = "chain_exact_size", since = "1.41.0")]
impl<A, B> ExactSizeIterator for Chain<A, B>
where
A: ExactSizeIterator,
B: ExactSizeIterator<Item = A::Item>,
{
/// Returns the exact number of times the iterator will iterate.
///
/// # Overflow Behavior
///
/// Calling this method on an iterator with more than [`usize::MAX`]
/// elements will result in a panic.
///
/// # Panics
///
/// This panics if the iterator has more than [`usize::MAX`] elements.
///
/// [`usize::MAX`]: ../../std/usize/constant.MAX.html
fn len(&self) -> usize {
match self.state {
ChainState::Both => {
self.a.len().checked_add(self.b.len())
.expect("called `len` on iterator with more than `usize::MAX` elements")
}
ChainState::Front => self.a.len(),
ChainState::Back => self.b.len(),
}
}
}

#[unstable(feature = "trusted_len", issue = "37572")]
unsafe impl<A, B> TrustedLen for Chain<A, B>
where A: TrustedLen, B: TrustedLen<Item=A::Item>,
Expand Down
24 changes: 24 additions & 0 deletions src/libcore/iter/traits/iterator.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
// ignore-tidy-filelength
// This file almost exclusively consists of the definition of `Iterator`. We
// can't split that into multiple files.

use crate::cmp::{self, Ordering};
use crate::ops::{Add, Try};

Expand Down Expand Up @@ -388,6 +392,26 @@ pub trait Iterator {
/// [`once`] is commonly used to adapt a single value into a chain of
/// other kinds of iteration.
///
///
/// # Overflowing behavior for long iterators
///
/// This method allows to easily build an iterator that yields more than
/// [`usize::MAX`] items. In that case, some methods that return `usize`
/// are not guarded against overflow. For example, this includes the
/// following methods:
///
/// - [`Iterator::count`]
/// - [`Iterator::enumerate`]
/// - [`Iterator::position`] and [`Iterator::rposition`]
/// - [`ExactSizeIterator::len`]
///
/// An overflow in those methods leads to a wrong result or a panic. If
/// debug assertions are enabled, a panic is guaranteed.
///
///
/// [`usize::MAX`]: ../../std/usize/constant.MAX.html
///
///
/// # Examples
///
/// Basic usage:
Expand Down
44 changes: 44 additions & 0 deletions src/libcore/tests/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,50 @@ fn test_iterator_chain_size_hint() {
assert_eq!(iter.size_hint(), (0, Some(0)));
}

#[test]
fn test_iterator_chain_len() {
let xs = [0, 1, 2];
let ys = [30, 40, 50, 60];

// First iterator is exhausted first
let mut iter = xs.iter().chain(&ys);
assert_eq!(iter.len(), 7);
assert_eq!(iter.next(), Some(&0));
assert_eq!(iter.len(), 6);
assert_eq!(iter.next(), Some(&1));
assert_eq!(iter.len(), 5);
assert_eq!(iter.next_back(), Some(&60));
assert_eq!(iter.len(), 4);
assert_eq!(iter.next(), Some(&2));
assert_eq!(iter.len(), 3);
assert_eq!(iter.next(), Some(&30));
assert_eq!(iter.len(), 2);
assert_eq!(iter.next(), Some(&40));
assert_eq!(iter.len(), 1);
assert_eq!(iter.next(), Some(&50));
assert_eq!(iter.len(), 0);
assert_eq!(iter.next(), None);

// Second iterator is exhausted first
let mut iter = xs.iter().chain(&ys);
assert_eq!(iter.len(), 7);
assert_eq!(iter.next_back(), Some(&60));
assert_eq!(iter.len(), 6);
assert_eq!(iter.next(), Some(&0));
assert_eq!(iter.len(), 5);
assert_eq!(iter.next_back(), Some(&50));
assert_eq!(iter.len(), 4);
assert_eq!(iter.next_back(), Some(&40));
assert_eq!(iter.len(), 3);
assert_eq!(iter.next_back(), Some(&30));
assert_eq!(iter.len(), 2);
assert_eq!(iter.next(), Some(&1));
assert_eq!(iter.len(), 1);
assert_eq!(iter.next(), Some(&2));
assert_eq!(iter.len(), 0);
assert_eq!(iter.next(), None);
}

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

0 comments on commit 722bba0

Please sign in to comment.