1 use core::iter::{Skip, Take};
2 use core::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};
3
4 #[cfg(doc)]
5 use crate::Itertools;
6
7 mod 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.
22 pub trait IteratorIndex<I>: private_iter_index::Sealed
23 where
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.
index(self, from: I) -> Self::Output33 fn index(self, from: I) -> Self::Output;
34 }
35
36 impl<I> IteratorIndex<I> for Range<usize>
37 where
38 I: Iterator,
39 {
40 type Output = Skip<Take<I>>;
41
index(self, iter: I) -> Self::Output42 fn index(self, iter: I) -> Self::Output {
43 iter.take(self.end).skip(self.start)
44 }
45 }
46
47 impl<I> IteratorIndex<I> for RangeInclusive<usize>
48 where
49 I: Iterator,
50 {
51 type Output = Take<Skip<I>>;
52
index(self, iter: I) -> Self::Output53 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
65 impl<I> IteratorIndex<I> for RangeTo<usize>
66 where
67 I: Iterator,
68 {
69 type Output = Take<I>;
70
index(self, iter: I) -> Self::Output71 fn index(self, iter: I) -> Self::Output {
72 iter.take(self.end)
73 }
74 }
75
76 impl<I> IteratorIndex<I> for RangeToInclusive<usize>
77 where
78 I: Iterator,
79 {
80 type Output = Take<I>;
81
index(self, iter: I) -> Self::Output82 fn index(self, iter: I) -> Self::Output {
83 assert_ne!(self.end, usize::MAX);
84 iter.take(self.end + 1)
85 }
86 }
87
88 impl<I> IteratorIndex<I> for RangeFrom<usize>
89 where
90 I: Iterator,
91 {
92 type Output = Skip<I>;
93
index(self, iter: I) -> Self::Output94 fn index(self, iter: I) -> Self::Output {
95 iter.skip(self.start)
96 }
97 }
98
99 impl<I> IteratorIndex<I> for RangeFull
100 where
101 I: Iterator,
102 {
103 type Output = I;
104
index(self, iter: I) -> Self::Output105 fn index(self, iter: I) -> Self::Output {
106 iter
107 }
108 }
109
get<I, R>(iter: I, index: R) -> R::Output where I: IntoIterator, R: IteratorIndex<I::IntoIter>,110 pub fn get<I, R>(iter: I, index: R) -> R::Output
111 where
112 I: IntoIterator,
113 R: IteratorIndex<I::IntoIter>,
114 {
115 index.index(iter.into_iter())
116 }
117