From 3f9b26dc64a2068d30027fd29ffbbfe07663419f Mon Sep 17 00:00:00 2001
From: The8472 <git@infinite-source.de>
Date: Thu, 14 Oct 2021 12:10:18 +0200
Subject: [PATCH] Fix Iterator::advance_by contract inconsistency

The `advance_by(n)` docs state that in the error case `Err(k)` that k is always less than n.
It also states that `advance_by(0)` may return `Err(0)` to indicate an exhausted iterator.
These statements are inconsistent.
Since only one implementation (Skip) actually made use of that I changed it to return Ok(()) in that case too.

While adding some tests I also found a bug in `Take::advance_back_by`.
---
 library/alloc/tests/vec.rs                   |  3 ++
 library/core/src/iter/adapters/skip.rs       |  4 +--
 library/core/src/iter/adapters/take.rs       | 29 ++++++++++----------
 library/core/src/iter/traits/double_ended.rs |  3 --
 library/core/src/iter/traits/iterator.rs     |  3 --
 library/core/tests/iter/adapters/chain.rs    |  8 ++++++
 library/core/tests/iter/adapters/flatten.rs  |  3 ++
 library/core/tests/iter/adapters/skip.rs     | 11 ++++++++
 library/core/tests/iter/adapters/take.rs     | 22 +++++++++++++++
 library/core/tests/iter/range.rs             |  3 ++
 library/core/tests/slice.rs                  |  2 ++
 11 files changed, 69 insertions(+), 22 deletions(-)

diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs
index 00a878c079480..773142825329b 100644
--- a/library/alloc/tests/vec.rs
+++ b/library/alloc/tests/vec.rs
@@ -985,6 +985,9 @@ fn test_into_iter_advance_by() {
 
     assert_eq!(i.advance_by(usize::MAX), Err(0));
 
+    i.advance_by(0).unwrap();
+    i.advance_back_by(0).unwrap();
+
     assert_eq!(i.len(), 0);
 }
 
diff --git a/library/core/src/iter/adapters/skip.rs b/library/core/src/iter/adapters/skip.rs
index 565fc224f53ca..ea1da8ba434ed 100644
--- a/library/core/src/iter/adapters/skip.rs
+++ b/library/core/src/iter/adapters/skip.rs
@@ -119,8 +119,8 @@ where
     #[rustc_inherit_overflow_checks]
     fn advance_by(&mut self, n: usize) -> Result<(), usize> {
         let mut rem = n;
-
         let step_one = self.n.saturating_add(rem);
+
         match self.iter.advance_by(step_one) {
             Ok(_) => {
                 rem -= step_one - self.n;
@@ -129,7 +129,7 @@ where
             Err(advanced) => {
                 let advanced_without_skip = advanced.saturating_sub(self.n);
                 self.n = self.n.saturating_sub(advanced);
-                return Err(advanced_without_skip);
+                return if n == 0 { Ok(()) } else { Err(advanced_without_skip) };
             }
         }
 
diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs
index 81f6c294fac13..2962e0104d11d 100644
--- a/library/core/src/iter/adapters/take.rs
+++ b/library/core/src/iter/adapters/take.rs
@@ -215,21 +215,22 @@ where
     }
 
     #[inline]
+    #[rustc_inherit_overflow_checks]
     fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
-        let inner_len = self.iter.len();
-        let len = self.n;
-        let remainder = len.saturating_sub(n);
-        let to_advance = inner_len - remainder;
-        match self.iter.advance_back_by(to_advance) {
-            Ok(_) => {
-                self.n = remainder;
-                if n > len {
-                    return Err(len);
-                }
-                return Ok(());
-            }
-            _ => panic!("ExactSizeIterator contract violation"),
-        }
+        // The amount by which the inner iterator needs to be shortened for it to be
+        // at most as long as the take() amount.
+        let trim_inner = self.iter.len().saturating_sub(self.n);
+        // The amount we need to advance inner to fulfill the caller's request.
+        // take(), advance_by() and len() all can be at most usize, so we don't have to worry
+        // about having to advance more than usize::MAX here.
+        let advance_by = trim_inner.saturating_add(n);
+
+        let advanced = match self.iter.advance_back_by(advance_by) {
+            Ok(_) => advance_by - trim_inner,
+            Err(advanced) => advanced - trim_inner,
+        };
+        self.n -= advanced;
+        return if advanced < n { Err(advanced) } else { Ok(()) };
     }
 }
 
diff --git a/library/core/src/iter/traits/double_ended.rs b/library/core/src/iter/traits/double_ended.rs
index 9a589c1f3b55c..a6aed6d210beb 100644
--- a/library/core/src/iter/traits/double_ended.rs
+++ b/library/core/src/iter/traits/double_ended.rs
@@ -106,9 +106,6 @@ pub trait DoubleEndedIterator: Iterator {
     /// Calling `advance_back_by(0)` can do meaningful work, for example [`Flatten`] can advance its
     /// outer iterator until it finds an inner iterator that is not empty, which then often
     /// allows it to return a more accurate `size_hint()` than in its initial state.
-    /// `advance_back_by(0)` may either return `Ok()` or `Err(0)`. The former conveys no information
-    /// whether the iterator is or is not exhausted, the latter can be treated as if [`next_back`]
-    /// had returned `None`. Replacing a `Err(0)` with `Ok` is only correct for `n = 0`.
     ///
     /// [`advance_by`]: Iterator::advance_by
     /// [`Flatten`]: crate::iter::Flatten
diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs
index d957a7527cf58..588758fed9cab 100644
--- a/library/core/src/iter/traits/iterator.rs
+++ b/library/core/src/iter/traits/iterator.rs
@@ -249,9 +249,6 @@ pub trait Iterator {
     /// Calling `advance_by(0)` can do meaningful work, for example [`Flatten`]
     /// can advance its outer iterator until it finds an inner iterator that is not empty, which
     /// then often allows it to return a more accurate `size_hint()` than in its initial state.
-    /// `advance_by(0)` may either return `Ok()` or `Err(0)`. The former conveys no information
-    /// whether the iterator is or is not exhausted, the latter can be treated as if [`next`]
-    /// had returned `None`. Replacing a `Err(0)` with `Ok` is only correct for `n = 0`.
     ///
     /// [`Flatten`]: crate::iter::Flatten
     /// [`next`]: Iterator::next
diff --git a/library/core/tests/iter/adapters/chain.rs b/library/core/tests/iter/adapters/chain.rs
index 4cd79687b53d4..f419f9cec12f8 100644
--- a/library/core/tests/iter/adapters/chain.rs
+++ b/library/core/tests/iter/adapters/chain.rs
@@ -34,6 +34,7 @@ fn test_iterator_chain_advance_by() {
             iter.advance_by(i).unwrap();
             assert_eq!(iter.next(), Some(&xs[i]));
             assert_eq!(iter.advance_by(100), Err(len - i - 1));
+            iter.advance_by(0).unwrap();
         }
 
         for i in 0..ys.len() {
@@ -41,14 +42,17 @@ fn test_iterator_chain_advance_by() {
             iter.advance_by(xs.len() + i).unwrap();
             assert_eq!(iter.next(), Some(&ys[i]));
             assert_eq!(iter.advance_by(100), Err(ys.len() - i - 1));
+            iter.advance_by(0).unwrap();
         }
 
         let mut iter = xs.iter().chain(ys);
         iter.advance_by(len).unwrap();
         assert_eq!(iter.next(), None);
+        iter.advance_by(0).unwrap();
 
         let mut iter = xs.iter().chain(ys);
         assert_eq!(iter.advance_by(len + 1), Err(len));
+        iter.advance_by(0).unwrap();
     }
 
     test_chain(&[], &[]);
@@ -67,6 +71,7 @@ fn test_iterator_chain_advance_back_by() {
             iter.advance_back_by(i).unwrap();
             assert_eq!(iter.next_back(), Some(&ys[ys.len() - i - 1]));
             assert_eq!(iter.advance_back_by(100), Err(len - i - 1));
+            iter.advance_back_by(0).unwrap();
         }
 
         for i in 0..xs.len() {
@@ -74,14 +79,17 @@ fn test_iterator_chain_advance_back_by() {
             iter.advance_back_by(ys.len() + i).unwrap();
             assert_eq!(iter.next_back(), Some(&xs[xs.len() - i - 1]));
             assert_eq!(iter.advance_back_by(100), Err(xs.len() - i - 1));
+            iter.advance_back_by(0).unwrap();
         }
 
         let mut iter = xs.iter().chain(ys);
         iter.advance_back_by(len).unwrap();
         assert_eq!(iter.next_back(), None);
+        iter.advance_back_by(0).unwrap();
 
         let mut iter = xs.iter().chain(ys);
         assert_eq!(iter.advance_back_by(len + 1), Err(len));
+        iter.advance_back_by(0).unwrap();
     }
 
     test_chain(&[], &[]);
diff --git a/library/core/tests/iter/adapters/flatten.rs b/library/core/tests/iter/adapters/flatten.rs
index 4ae50a2f06681..cd6513327f00e 100644
--- a/library/core/tests/iter/adapters/flatten.rs
+++ b/library/core/tests/iter/adapters/flatten.rs
@@ -61,6 +61,7 @@ fn test_flatten_try_folds() {
 #[test]
 fn test_flatten_advance_by() {
     let mut it = once(0..10).chain(once(10..30)).chain(once(30..40)).flatten();
+
     it.advance_by(5).unwrap();
     assert_eq!(it.next(), Some(5));
     it.advance_by(9).unwrap();
@@ -72,6 +73,8 @@ fn test_flatten_advance_by() {
 
     assert_eq!(it.advance_by(usize::MAX), Err(9));
     assert_eq!(it.advance_back_by(usize::MAX), Err(0));
+    it.advance_by(0).unwrap();
+    it.advance_back_by(0).unwrap();
     assert_eq!(it.size_hint(), (0, Some(0)));
 }
 
diff --git a/library/core/tests/iter/adapters/skip.rs b/library/core/tests/iter/adapters/skip.rs
index cf60057a1644a..0c464bdd03a22 100644
--- a/library/core/tests/iter/adapters/skip.rs
+++ b/library/core/tests/iter/adapters/skip.rs
@@ -69,6 +69,17 @@ fn test_iterator_skip_nth() {
     assert_eq!(it.nth(0), None);
 }
 
+#[test]
+fn test_skip_advance_by() {
+    assert_eq!((0..0).skip(10).advance_by(0), Ok(()));
+    assert_eq!((0..0).skip(10).advance_by(1), Err(0));
+    assert_eq!((0u128..(usize::MAX as u128) + 1).skip(usize::MAX).advance_by(usize::MAX), Err(1));
+    assert_eq!((0u128..u128::MAX).skip(usize::MAX).advance_by(1), Ok(()));
+
+    assert_eq!((0..2).skip(1).advance_back_by(10), Err(1));
+    assert_eq!((0..0).skip(1).advance_back_by(0), Ok(()));
+}
+
 #[test]
 fn test_iterator_skip_count() {
     let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
diff --git a/library/core/tests/iter/adapters/take.rs b/library/core/tests/iter/adapters/take.rs
index 89f9cb1e2ed1b..bfb659f0a8378 100644
--- a/library/core/tests/iter/adapters/take.rs
+++ b/library/core/tests/iter/adapters/take.rs
@@ -73,6 +73,28 @@ fn test_iterator_take_nth_back() {
     assert_eq!(it.nth_back(1), None);
 }
 
+#[test]
+fn test_take_advance_by() {
+    let mut take = (0..10).take(3);
+    assert_eq!(take.advance_by(2), Ok(()));
+    assert_eq!(take.next(), Some(2));
+    assert_eq!(take.advance_by(1), Err(0));
+
+    assert_eq!((0..0).take(10).advance_by(0), Ok(()));
+    assert_eq!((0..0).take(10).advance_by(1), Err(0));
+    assert_eq!((0..10).take(4).advance_by(5), Err(4));
+
+    let mut take = (0..10).take(3);
+    assert_eq!(take.advance_back_by(2), Ok(()));
+    assert_eq!(take.next(), Some(0));
+    assert_eq!(take.advance_back_by(1), Err(0));
+
+    assert_eq!((0..2).take(1).advance_back_by(10), Err(1));
+    assert_eq!((0..0).take(1).advance_back_by(1), Err(0));
+    assert_eq!((0..0).take(1).advance_back_by(0), Ok(()));
+    assert_eq!((0..usize::MAX).take(100).advance_back_by(usize::MAX), Err(100));
+}
+
 #[test]
 fn test_iterator_take_short() {
     let xs = [0, 1, 2, 3];
diff --git a/library/core/tests/iter/range.rs b/library/core/tests/iter/range.rs
index 6b4cf33efe1ff..84498a8eae52e 100644
--- a/library/core/tests/iter/range.rs
+++ b/library/core/tests/iter/range.rs
@@ -300,6 +300,9 @@ fn test_range_advance_by() {
 
     assert_eq!(r.advance_by(usize::MAX), Err(usize::MAX - 2));
 
+    r.advance_by(0).unwrap();
+    r.advance_back_by(0).unwrap();
+
     let mut r = 0u128..u128::MAX;
 
     r.advance_by(usize::MAX).unwrap();
diff --git a/library/core/tests/slice.rs b/library/core/tests/slice.rs
index 8d05e47edf489..66f03c68cbe01 100644
--- a/library/core/tests/slice.rs
+++ b/library/core/tests/slice.rs
@@ -154,6 +154,7 @@ fn test_iterator_advance_by() {
     assert_eq!(iter.as_slice(), &v[3..]);
     iter.advance_by(2).unwrap();
     assert_eq!(iter.as_slice(), &[]);
+    iter.advance_by(0).unwrap();
 }
 
 #[test]
@@ -175,6 +176,7 @@ fn test_iterator_advance_back_by() {
     assert_eq!(iter.as_slice(), &v[..v.len() - 3]);
     iter.advance_back_by(2).unwrap();
     assert_eq!(iter.as_slice(), &[]);
+    iter.advance_back_by(0).unwrap();
 }
 
 #[test]