Do not panic in Itertools::get large range

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<usize> {}
     impl Sealed for ops::RangeFrom<usize> {}
     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<usize>,
+    }
+
+    impl UsizeMaxPlusOne {
+        fn checked_dec(self) -> Option<Self> {
+            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<usize> for UsizeMaxPlusOne {
+        fn from(n: usize) -> Self {
+            UsizeMaxPlusOne { n: Some(n) }
+        }
+    }
+
+    impl PartialOrd for UsizeMaxPlusOne {
+        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+            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<I> {
+        pub(crate) iter: I,
+        pub(crate) rem: UsizeMaxPlusOne,
+    }
+
+    impl<I: Iterator> Iterator for TakeWithUsizeMaxPlusOne<I> {
+        type Item = I::Item;
+
+        fn next(&mut self) -> Option<Self::Item> {
+            self.rem = self.rem.checked_dec()?;
+            self.iter.next()
+        }
+
+        fn size_hint(&self) -> (usize, Option<usize>) {
+            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<I: DoubleEndedIterator> DoubleEndedIterator for TakeWithUsizeMaxPlusOne<I> {
+        fn next_back(&mut self) -> Option<Self::Item> {
+            self.rem = self.rem.checked_dec()?;
+            self.iter.next_back()
+        }
+    }
 }
 
 /// Used by [`Itertools::get`] to know which iterator
@@ -48,17 +134,14 @@ impl<I> IteratorIndex<I> for RangeInclusive<usize>
 where
     I: Iterator,
 {
-    type Output = Take<Skip<I>>;
+    type Output = TakeWithUsizeMaxPlusOne<Skip<I>>;
 
     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<I> IteratorIndex<I> for RangeToInclusive<usize>
 where
     I: Iterator,
 {
-    type Output = Take<I>;
+    type Output = TakeWithUsizeMaxPlusOne<I>;
 
     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,
+        }
     }
 }
 

tests/test_core.rs

@@ -22,21 +22,17 @@ use itertools as it;
 fn get_esi_then_esi<I: ExactSizeIterator + Clone>(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(..));
 }
 
 #[allow(dead_code)]
 fn get_dei_esi_then_dei_esi<I: DoubleEndedIterator + ExactSizeIterator + Clone>(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]