From 37896863e3c44f219c051a6c81df0da6300158c5 Mon Sep 17 00:00:00 2001 From: Stepan Koltsov Date: Mon, 23 Mar 2026 07:05:57 +0000 Subject: [PATCH] Do not panic in Itertools::get large range --- src/iter_index.rs | 115 +++++++++++++++++++++++++++++++++++++++------ tests/test_core.rs | 18 ++++--- 2 files changed, 112 insertions(+), 21 deletions(-) diff --git a/src/iter_index.rs b/src/iter_index.rs index aadaa72a7..843b0ca5f 100644 --- a/src/iter_index.rs +++ b/src/iter_index.rs @@ -1,10 +1,12 @@ -use core::iter::{Skip, Take}; -use core::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive}; - +use crate::iter_index::private_iter_index::{TakeWithUsizeMaxPlusOne, UsizeMaxPlusOne}; #[cfg(doc)] use crate::Itertools; +use core::iter::{Skip, Take}; +use core::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive}; mod private_iter_index { + use core::cmp; + use core::cmp::Ordering; use core::ops; pub trait Sealed {} @@ -15,6 +17,90 @@ mod private_iter_index { impl Sealed for ops::RangeToInclusive {} impl Sealed for ops::RangeFrom {} impl Sealed for ops::RangeFull {} + + /// An integer in range `0..=usize::MAX+1`. + #[derive(Copy, Clone, Eq, PartialEq, Debug)] + pub(crate) struct UsizeMaxPlusOne { + /// `None -> usize::MAX + 1`. + pub(crate) n: Option, + } + + impl UsizeMaxPlusOne { + fn checked_dec(self) -> Option { + Some(UsizeMaxPlusOne::from(match self.n { + None => usize::MAX, + Some(n) => n.checked_sub(1)?, + })) + } + + fn saturating_to_usize(self) -> usize { + self.n.unwrap_or(usize::MAX) + } + + pub(crate) fn inc(n: usize) -> Self { + UsizeMaxPlusOne { + n: n.checked_add(1), + } + } + } + + impl From for UsizeMaxPlusOne { + fn from(n: usize) -> Self { + UsizeMaxPlusOne { n: Some(n) } + } + } + + impl PartialOrd for UsizeMaxPlusOne { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + impl Ord for UsizeMaxPlusOne { + fn cmp(&self, other: &Self) -> Ordering { + match (self.n, other.n) { + (None, None) => Ordering::Equal, + (None, Some(_)) => Ordering::Greater, + (Some(_), None) => Ordering::Less, + (Some(self_n), Some(other_n)) => self_n.cmp(&other_n), + } + } + } + + #[derive(Debug)] + pub struct TakeWithUsizeMaxPlusOne { + pub(crate) iter: I, + pub(crate) rem: UsizeMaxPlusOne, + } + + impl Iterator for TakeWithUsizeMaxPlusOne { + type Item = I::Item; + + fn next(&mut self) -> Option { + self.rem = self.rem.checked_dec()?; + self.iter.next() + } + + fn size_hint(&self) -> (usize, Option) { + let (lower, upper) = self.iter.size_hint(); + let lower = UsizeMaxPlusOne::from(lower); + let upper = upper.map(UsizeMaxPlusOne::from); + + let ret_lower = cmp::min(self.rem, lower).saturating_to_usize(); + let ret_upper = match upper { + None => self.rem.n, + Some(upper) => cmp::min(self.rem, upper).n, + }; + (ret_lower, ret_upper) + } + } + + impl DoubleEndedIterator for TakeWithUsizeMaxPlusOne { + fn next_back(&mut self) -> Option { + self.rem = self.rem.checked_dec()?; + self.iter.next_back() + } + } } /// Used by [`Itertools::get`] to know which iterator @@ -48,17 +134,14 @@ impl IteratorIndex for RangeInclusive where I: Iterator, { - type Output = Take>; + type Output = TakeWithUsizeMaxPlusOne>; fn index(self, iter: I) -> Self::Output { - // end - start + 1 without overflowing if possible - let length = if *self.end() == usize::MAX { - assert_ne!(*self.start(), 0); - self.end() - self.start() + 1 - } else { - (self.end() + 1).saturating_sub(*self.start()) - }; - iter.skip(*self.start()).take(length) + let length = UsizeMaxPlusOne::inc(self.end().saturating_sub(*self.start())); + TakeWithUsizeMaxPlusOne { + rem: length, + iter: iter.skip(*self.start()), + } } } @@ -77,11 +160,13 @@ impl IteratorIndex for RangeToInclusive where I: Iterator, { - type Output = Take; + type Output = TakeWithUsizeMaxPlusOne; fn index(self, iter: I) -> Self::Output { - assert_ne!(self.end, usize::MAX); - iter.take(self.end + 1) + TakeWithUsizeMaxPlusOne { + rem: UsizeMaxPlusOne::inc(self.end), + iter, + } } } diff --git a/tests/test_core.rs b/tests/test_core.rs index cb2467cfb..cfb5c4c24 100644 --- a/tests/test_core.rs +++ b/tests/test_core.rs @@ -22,10 +22,8 @@ use itertools as it; fn get_esi_then_esi(it: I) { fn is_esi(_: impl ExactSizeIterator) {} is_esi(it.clone().get(1..4)); - is_esi(it.clone().get(1..=4)); is_esi(it.clone().get(1..)); is_esi(it.clone().get(..4)); - is_esi(it.clone().get(..=4)); is_esi(it.get(..)); } @@ -33,10 +31,8 @@ fn get_esi_then_esi(it: I) { fn get_dei_esi_then_dei_esi(it: I) { fn is_dei_esi(_: impl DoubleEndedIterator + ExactSizeIterator) {} is_dei_esi(it.clone().get(1..4)); - is_dei_esi(it.clone().get(1..=4)); is_dei_esi(it.clone().get(1..)); is_dei_esi(it.clone().get(..4)); - is_dei_esi(it.clone().get(..=4)); is_dei_esi(it.get(..)); } @@ -48,9 +44,19 @@ fn get_1_max() { } #[test] -#[should_panic] fn get_full_range_inclusive() { - let _it = (0..5).get(0..=usize::MAX); + let mut it = (0..5).get(0..=usize::MAX); + assert_eq!((5, Some(5)), it.size_hint()); + assert_eq!(it.next(), Some(0)); + assert_eq!(it.next_back(), Some(4)); + assert_eq!((3, Some(3)), it.size_hint()); + assert_eq!(it.next(), Some(1)); + assert_eq!(it.next(), Some(2)); + assert_eq!(it.next(), Some(3)); + assert_eq!(it.next(), None); + + let it = iter::repeat(0).get(0..=usize::MAX); + assert_eq!((usize::MAX, None), it.size_hint()); } #[test]