Use saturating arithmetic to avoid usize overflow panics (#1066, #1067, #1068)

src/adaptors/mod.rs

@@ -179,7 +179,7 @@ where
         let (combined_lower, combined_upper) =
             size_hint::mul_scalar(size_hint::min(curr_hint, next_hint), 2);
         let lower = if curr_lower > next_lower {
-            combined_lower + 1
+            combined_lower.saturating_add(1)
         } else {
             combined_lower
         };

src/k_smallest.rs

@@ -99,7 +99,10 @@ where
     }
 
     let mut iter = iter.fuse();
-    let mut buf = iter.by_ref().take(2 * k).collect::<Vec<_>>();
+    let mut buf = iter
+        .by_ref()
+        .take(2usize.saturating_mul(k))
+        .collect::<Vec<_>>();
 
     if buf.len() < k {
         buf.sort_unstable_by(&mut comparator);

src/peek_nth.rs

@@ -69,7 +69,7 @@ where
     /// assert_eq!(iter.peek_nth(1), None);
     /// ```
     pub fn peek_nth(&mut self, n: usize) -> Option<&I::Item> {
-        let unbuffered_items = (n + 1).saturating_sub(self.buf.len());
+        let unbuffered_items = n.saturating_add(1).saturating_sub(self.buf.len());
 
         self.buf.extend(self.iter.by_ref().take(unbuffered_items));
 
@@ -110,7 +110,7 @@ where
     /// assert_eq!(iter.peek_nth_mut(1), None);
     /// ```
     pub fn peek_nth_mut(&mut self, n: usize) -> Option<&mut I::Item> {
-        let unbuffered_items = (n + 1).saturating_sub(self.buf.len());
+        let unbuffered_items = n.saturating_add(1).saturating_sub(self.buf.len());
 
         self.buf.extend(self.iter.by_ref().take(unbuffered_items));
 

tests/test_std.rs

@@ -86,6 +86,16 @@ fn interleave_shortest() {
     assert_eq!(it.size_hint(), (6, Some(6)));
 }
 
+#[test]
+fn interleave_shortest_size_hint_does_not_overflow() {
+    // The combined lower bound can exceed `usize::MAX`, which used to overflow
+    // when computing `size_hint`. It should saturate instead of panicking.
+    let i = ::std::iter::repeat(0u8).take(usize::MAX);
+    let j = ::std::iter::repeat(0u8).take(usize::MAX - 1);
+    let it = i.interleave_shortest(j);
+    assert_eq!(it.size_hint(), (usize::MAX, None));
+}
+
 #[test]
 fn duplicates_by() {
     let xs = ["aaa", "bbbbb", "aa", "ccc", "bbbb", "aaaaa", "cccc"];
@@ -532,6 +542,21 @@ fn sorted_by() {
     it::assert_equal(v, vec![4, 3, 2, 1, 0]);
 }
 
+#[test]
+fn k_smallest_relaxed_does_not_overflow() {
+    // `k_smallest_relaxed` allocates room for `2 * k` items. A large `k` used to
+    // overflow that multiplication and panic; it should saturate instead.
+    it::assert_equal(
+        vec![1].into_iter().k_smallest_relaxed(usize::MAX / 2),
+        vec![1],
+    );
+    it::assert_equal(
+        vec![1].into_iter().k_smallest_relaxed(usize::MAX / 2 + 1),
+        vec![1],
+    );
+    it::assert_equal(vec![1].into_iter().k_smallest_relaxed(usize::MAX), vec![1]);
+}
+
 #[cfg(not(miri))]
 qc::quickcheck! {
     fn k_smallest_range(n: i64, m: u16, k: u16) -> () {
@@ -835,6 +860,20 @@ fn test_peek_nth() {
     assert_eq!(iter.peek_nth(1), None);
 }
 
+#[test]
+fn test_peek_nth_max_does_not_overflow() {
+    // Peeking past the end returns `None`. Asking for `usize::MAX` used to
+    // overflow when computing `n + 1`; it should return `None` instead.
+    let nums = vec![1u8, 2, 3];
+
+    let mut iter = peek_nth(nums.iter().copied());
+    assert_eq!(iter.peek_nth(usize::MAX), None);
+    assert_eq!(iter.peek_nth_mut(usize::MAX), None);
+
+    // The iterator is still intact and yields all of its items.
+    assert_eq!(iter.collect::<Vec<_>>(), nums);
+}
+
 #[test]
 fn test_peek_nth_peeking_next() {
     use it::PeekingNext;