From 9d3e84432dae2e96a5e0f97be18ee09b5a2217b1 Mon Sep 17 00:00:00 2001
From: Oliver Middleton <olliemail27@gmail.com>
Date: Wed, 22 Jan 2020 20:28:28 +0000
Subject: [PATCH] Avoid overflow in `std::iter::Skip::count`

The call to `count` on the inner iterator can overflow even if `Skip` itself would return less that `usize::max_value()` items.
---
 src/libcore/iter/adapters/mod.rs             | 10 ++++++++--
 src/test/ui/iterators/skip-count-overflow.rs |  8 ++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)
 create mode 100644 src/test/ui/iterators/skip-count-overflow.rs

diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs
index 6eb837ed0fed8..5787b9174edab 100644
--- a/src/libcore/iter/adapters/mod.rs
+++ b/src/libcore/iter/adapters/mod.rs
@@ -1815,8 +1815,14 @@ where
     }
 
     #[inline]
-    fn count(self) -> usize {
-        self.iter.count().saturating_sub(self.n)
+    fn count(mut self) -> usize {
+        if self.n > 0 {
+            // nth(n) skips n+1
+            if self.iter.nth(self.n - 1).is_none() {
+                return 0;
+            }
+        }
+        self.iter.count()
     }
 
     #[inline]
diff --git a/src/test/ui/iterators/skip-count-overflow.rs b/src/test/ui/iterators/skip-count-overflow.rs
new file mode 100644
index 0000000000000..d8efc948664ff
--- /dev/null
+++ b/src/test/ui/iterators/skip-count-overflow.rs
@@ -0,0 +1,8 @@
+// run-pass
+// only-32bit too impatient for 2⁶⁴ items
+// compile-flags: -C overflow-checks -C opt-level=3
+
+fn main() {
+    let i = (0..usize::max_value()).chain(0..10).skip(usize::max_value());
+    assert_eq!(i.count(), 10);
+}