Skip to main content

enumerate_number/
lib.rs

1#![doc = include_str!("../README.md")]
2#![no_std]
3use core::iter::FusedIterator;
4
5#[doc(hidden)]
6pub trait Counter: Copy + Default {
7    fn inc(&mut self);
8    fn dec(&mut self);
9    fn inc_n(&mut self, n: usize);
10}
11macro_rules! impl_counter {
12    ($ty:ty) => {
13        impl Counter for $ty {
14            #[inline]
15            fn inc(&mut self) { *self += 1 as $ty }
16
17            #[inline]
18            fn dec(&mut self) { *self -= 1 as $ty }
19
20            #[inline]
21            fn inc_n(&mut self, n: usize) { *self += n as $ty }
22        }
23    };
24}
25impl_counter!(i8);
26impl_counter!(i16);
27impl_counter!(i32);
28impl_counter!(i64);
29impl_counter!(i128);
30impl_counter!(isize);
31impl_counter!(u8);
32impl_counter!(u16);
33impl_counter!(u32);
34impl_counter!(u64);
35impl_counter!(u128);
36impl_counter!(usize);
37impl_counter!(f32);
38impl_counter!(f64);
39
40#[derive(Debug, Clone, Default)]
41pub struct Enumerate<I: Iterator, C: Counter> {
42    iter: I,
43    count: C,
44}
45
46impl<I: Iterator, C: Counter> Iterator for Enumerate<I, C> {
47    type Item = (C, I::Item);
48
49    #[inline]
50    fn next(&mut self) -> Option<Self::Item> {
51        let a = self.iter.next()?;
52        let i = self.count;
53        self.count.inc();
54        Some((i, a))
55    }
56
57    #[inline]
58    fn size_hint(&self) -> (usize, Option<usize>) {
59        self.iter.size_hint()
60    }
61
62    #[inline]
63    fn nth(&mut self, n: usize) -> Option<Self::Item> {
64        let a = self.iter.nth(n)?;
65        self.count.inc_n(n);
66        let i = self.count;
67        self.count.inc();
68        Some((i, a))
69    }
70
71    #[inline]
72    fn count(self) -> usize {
73        self.iter.count()
74    }
75
76    #[inline]
77    fn fold<B, F>(self, init: B, mut f: F) -> B
78    where F: FnMut(B, Self::Item) -> B,
79    {
80        let mut count = self.count;
81        self.iter.fold(init, |acc, ele| {
82            let acc = f(acc, (count, ele));
83            count.inc();
84            acc
85        })
86    }
87}
88
89impl<I, C> DoubleEndedIterator for Enumerate<I, C>
90where I: DoubleEndedIterator + ExactSizeIterator,
91      C: Counter,
92{
93    #[inline]
94    fn next_back(&mut self) -> Option<Self::Item> {
95        let a = self.iter.next_back()?;
96        let len = self.iter.len();
97        let mut count = self.count;
98        count.inc_n(len);
99        Some((count, a))
100    }
101
102    #[inline]
103    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
104        let a = self.iter.nth_back(n)?;
105        let len = self.iter.len();
106        let mut count = self.count;
107        count.inc_n(len);
108        Some((count, a))
109    }
110
111    fn rfold<B, F>(self, init: B, mut f: F) -> B
112    where F: FnMut(B, Self::Item) -> B,
113    {
114        let mut count = self.count;
115        count.inc_n(self.iter.len());
116        self.iter.rfold(init, |acc, ele| {
117            count.dec();
118            f(acc, (count, ele))
119        })
120    }
121}
122
123impl<I: FusedIterator, C: Counter> FusedIterator for Enumerate<I, C> {}
124
125impl<I: ExactSizeIterator, C: Counter> ExactSizeIterator for Enumerate<I, C> {
126    fn len(&self) -> usize {
127        self.iter.len()
128    }
129}
130
131macro_rules! def_iterator_ext {
132    ($name:ident : $ty:ty) => {
133        /// Like [`EnumerateNumber::enumerate_number`]
134        #[inline]
135        fn $name(self) -> Enumerate<Self, $ty> {
136            Enumerate { iter: self, count: Default::default() }
137        }
138    };
139}
140
141pub trait EnumerateNumber: Iterator + Sized {
142    def_iterator_ext!(enumerate_i8: i8);
143    def_iterator_ext!(enumerate_i16: i16);
144    def_iterator_ext!(enumerate_i32: i32);
145    def_iterator_ext!(enumerate_i64: i64);
146    def_iterator_ext!(enumerate_i128: i128);
147    def_iterator_ext!(enumerate_isize: isize);
148    def_iterator_ext!(enumerate_u8: u8);
149    def_iterator_ext!(enumerate_u16: u16);
150    def_iterator_ext!(enumerate_u32: u32);
151    def_iterator_ext!(enumerate_u64: u64);
152    def_iterator_ext!(enumerate_u128: u128);
153    def_iterator_ext!(enumerate_f32: f32);
154    def_iterator_ext!(enumerate_f64: f64);
155
156    /// Use other number for enumerate
157    ///
158    /// # Examples
159    ///
160    /// ```
161    /// use enumerate_number::EnumerateNumber as _;
162    ///
163    /// let iter = "some".chars().enumerate_number();
164    /// let vec = iter.collect::<Vec<_>>();
165    /// assert_eq!(vec, vec![(0.0, 's'), (1.0, 'o'), (2.0, 'm'), (3.0, 'e')])
166    /// ```
167    #[inline]
168    fn enumerate_number<N: Counter>(self) -> Enumerate<Self, N> {
169        Enumerate { iter: self, count: Default::default() }
170    }
171}
172impl<I: Iterator> EnumerateNumber for I { }
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    extern crate alloc;
179    use alloc::vec;
180
181    #[test]
182    fn fold() {
183        let mut elems = vec![];
184        (0..5).enumerate_i16().fold((), |(), ele| {
185            elems.push(ele);
186        });
187        assert_eq!(elems, vec![
188            (0, 0),
189            (1, 1),
190            (2, 2),
191            (3, 3),
192            (4, 4),
193        ]);
194    }
195
196    #[test]
197    fn fold1() {
198        let mut elems = vec![];
199        let mut iter = (0..5).enumerate_i16();
200        assert_eq!(iter.next(), Some((0, 0)));
201        iter.fold((), |(), ele| {
202            elems.push(ele);
203        });
204        assert_eq!(elems, vec![
205            (1, 1),
206            (2, 2),
207            (3, 3),
208            (4, 4),
209        ]);
210    }
211
212    #[test]
213    fn rfold() {
214        let mut elems = vec![];
215        (0..5).enumerate_i16().rfold((), |(), ele| {
216            elems.push(ele);
217        });
218        assert_eq!(elems, vec![
219            (4, 4),
220            (3, 3),
221            (2, 2),
222            (1, 1),
223            (0, 0),
224        ]);
225    }
226
227    #[test]
228    fn rfold1() {
229        let mut elems = vec![];
230        let mut iter = (0..5).enumerate_i16();
231        assert_eq!(iter.next(), Some((0, 0)));
232        iter.rfold((), |(), ele| {
233            elems.push(ele);
234        });
235        assert_eq!(elems, vec![
236            (4, 4),
237            (3, 3),
238            (2, 2),
239            (1, 1),
240        ]);
241    }
242
243    #[test]
244    fn rfold2() {
245        let mut elems = vec![];
246        let mut iter = (0..5).enumerate_i16();
247        assert_eq!(iter.next_back(), Some((4, 4)));
248        iter.rfold((), |(), ele| {
249            elems.push(ele);
250        });
251        assert_eq!(elems, vec![
252            (3, 3),
253            (2, 2),
254            (1, 1),
255            (0, 0),
256        ]);
257    }
258
259    #[test]
260    fn nth() {
261        let mut iter = (0..5).enumerate_i16();
262        assert_eq!(iter.nth(1), Some((1, 1)));
263        assert_eq!(iter.nth(0), Some((2, 2)));
264        assert_eq!(iter.nth(1), Some((4, 4)));
265    }
266
267    #[test]
268    fn nth_back() {
269        let mut iter = (0..5).enumerate_i16();
270        assert_eq!(iter.nth_back(1), Some((3, 3)));
271        assert_eq!(iter.nth_back(0), Some((2, 2)));
272        assert_eq!(iter.nth_back(1), Some((0, 0)));
273    }
274
275    #[test]
276    fn nth_back1() {
277        let mut iter = (0..5).enumerate_i16();
278        assert_eq!(iter.nth_back(1), Some((3, 3)));
279        assert_eq!(iter.nth(0), Some((0, 0)));
280        assert_eq!(iter.nth(0), Some((1, 1)));
281    }
282}