diff --git a/src/adaptors/mod.rs b/src/adaptors/mod.rs index b5c0d42b1..451a4ee77 100644 --- a/src/adaptors/mod.rs +++ b/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 }; diff --git a/src/k_smallest.rs b/src/k_smallest.rs index 7e4ace262..c28831b28 100644 --- a/src/k_smallest.rs +++ b/src/k_smallest.rs @@ -99,7 +99,10 @@ where } let mut iter = iter.fuse(); - let mut buf = iter.by_ref().take(2 * k).collect::>(); + let mut buf = iter + .by_ref() + .take(2usize.saturating_mul(k)) + .collect::>(); if buf.len() < k { buf.sort_unstable_by(&mut comparator); diff --git a/src/peek_nth.rs b/src/peek_nth.rs index b03a3ef5f..4a176609d 100644 --- a/src/peek_nth.rs +++ b/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)); diff --git a/tests/test_std.rs b/tests/test_std.rs index 5f5cfd1eb..89a224df7 100644 --- a/tests/test_std.rs +++ b/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::>(), nums); +} + #[test] fn test_peek_nth_peeking_next() { use it::PeekingNext;