Add ExactSizeIterator impl for iter::Chain

Author: LukasKalbertodt

src/libcore/iter/adapters/chain.rs

@@ -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>,

src/libcore/iter/traits/iterator.rs

@@ -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};
 
@@ -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:

src/libcore/tests/iter.rs

@@ -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];