itertools/
iter_index.rs

1use core::iter::{Skip, Take};
2use core::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};
3
4#[cfg(doc)]
5use crate::Itertools;
6
7mod private_iter_index {
8    use core::ops;
9
10    pub trait Sealed {}
11
12    impl Sealed for ops::Range<usize> {}
13    impl Sealed for ops::RangeInclusive<usize> {}
14    impl Sealed for ops::RangeTo<usize> {}
15    impl Sealed for ops::RangeToInclusive<usize> {}
16    impl Sealed for ops::RangeFrom<usize> {}
17    impl Sealed for ops::RangeFull {}
18}
19
20/// Used by [`Itertools::get`] to know which iterator
21/// to turn different ranges into.
22pub trait IteratorIndex<I>: private_iter_index::Sealed
23where
24    I: Iterator,
25{
26    /// The type returned for this type of index.
27    type Output: Iterator<Item = I::Item>;
28
29    /// Returns an adapted iterator for the current index.
30    ///
31    /// Prefer calling [`Itertools::get`] instead
32    /// of calling this directly.
33    fn index(self, from: I) -> Self::Output;
34}
35
36impl<I> IteratorIndex<I> for Range<usize>
37where
38    I: Iterator,
39{
40    type Output = Skip<Take<I>>;
41
42    fn index(self, iter: I) -> Self::Output {
43        iter.take(self.end).skip(self.start)
44    }
45}
46
47impl<I> IteratorIndex<I> for RangeInclusive<usize>
48where
49    I: Iterator,
50{
51    type Output = Take<Skip<I>>;
52
53    fn index(self, iter: I) -> Self::Output {
54        // end - start + 1 without overflowing if possible
55        let length = if *self.end() == usize::MAX {
56            assert_ne!(*self.start(), 0);
57            self.end() - self.start() + 1
58        } else {
59            (self.end() + 1).saturating_sub(*self.start())
60        };
61        iter.skip(*self.start()).take(length)
62    }
63}
64
65impl<I> IteratorIndex<I> for RangeTo<usize>
66where
67    I: Iterator,
68{
69    type Output = Take<I>;
70
71    fn index(self, iter: I) -> Self::Output {
72        iter.take(self.end)
73    }
74}
75
76impl<I> IteratorIndex<I> for RangeToInclusive<usize>
77where
78    I: Iterator,
79{
80    type Output = Take<I>;
81
82    fn index(self, iter: I) -> Self::Output {
83        assert_ne!(self.end, usize::MAX);
84        iter.take(self.end + 1)
85    }
86}
87
88impl<I> IteratorIndex<I> for RangeFrom<usize>
89where
90    I: Iterator,
91{
92    type Output = Skip<I>;
93
94    fn index(self, iter: I) -> Self::Output {
95        iter.skip(self.start)
96    }
97}
98
99impl<I> IteratorIndex<I> for RangeFull
100where
101    I: Iterator,
102{
103    type Output = I;
104
105    fn index(self, iter: I) -> Self::Output {
106        iter
107    }
108}
109
110pub fn get<I, R>(iter: I, index: R) -> R::Output
111where
112    I: IntoIterator,
113    R: IteratorIndex<I::IntoIter>,
114{
115    index.index(iter.into_iter())
116}
itertools/iter_index.rs

itertools/
iter_index.rs
