polars_utils/itertools/
enumerate_idx.rs

1use num_traits::{FromPrimitive, One, Zero};
2
3/// An iterator that yields the current count and the element during iteration.
4///
5/// This `struct` is created by the [`enumerate`] method on [`Iterator`]. See its
6/// documentation for more.
7///
8/// [`enumerate`]: Iterator::enumerate
9/// [`Iterator`]: trait.Iterator.html
10#[derive(Clone, Debug)]
11#[must_use = "iterators are lazy and do nothing unless consumed"]
12pub struct EnumerateIdx<I, IdxType> {
13    iter: I,
14    count: IdxType,
15}
16
17impl<I, IdxType: Zero> EnumerateIdx<I, IdxType> {
18    pub fn new(iter: I) -> Self {
19        Self {
20            iter,
21            count: IdxType::zero(),
22        }
23    }
24}
25
26impl<I, IdxType> Iterator for EnumerateIdx<I, IdxType>
27where
28    I: Iterator,
29    IdxType: std::ops::Add<Output = IdxType> + FromPrimitive + std::ops::AddAssign + One + Copy,
30{
31    type Item = (IdxType, <I as Iterator>::Item);
32
33    /// # Overflow Behavior
34    ///
35    /// The method does no guarding against overflows, so enumerating more than
36    /// `idx::MAX` elements either produces the wrong result or panics. If
37    /// debug assertions are enabled, a panic is guaranteed.
38    ///
39    /// # Panics
40    ///
41    /// Might panic if the index of the element overflows a `idx`.
42    #[inline]
43    fn next(&mut self) -> Option<Self::Item> {
44        let a = self.iter.next()?;
45        let i = self.count;
46        self.count += IdxType::one();
47        Some((i, a))
48    }
49
50    #[inline]
51    fn size_hint(&self) -> (usize, Option<usize>) {
52        self.iter.size_hint()
53    }
54
55    #[inline]
56    fn nth(&mut self, n: usize) -> Option<Self::Item> {
57        let a = self.iter.nth(n)?;
58        let i = self.count + IdxType::from_usize(n).unwrap();
59        self.count = i + IdxType::one();
60        Some((i, a))
61    }
62
63    #[inline]
64    fn count(self) -> usize {
65        self.iter.count()
66    }
67}
68
69impl<I, IdxType> DoubleEndedIterator for EnumerateIdx<I, IdxType>
70where
71    I: ExactSizeIterator + DoubleEndedIterator,
72    IdxType: std::ops::Add<Output = IdxType> + FromPrimitive + std::ops::AddAssign + One + Copy,
73{
74    #[inline]
75    fn next_back(&mut self) -> Option<(IdxType, <I as Iterator>::Item)> {
76        let a = self.iter.next_back()?;
77        let len = IdxType::from_usize(self.iter.len()).unwrap();
78        // Can safely add, `ExactSizeIterator` promises that the number of
79        // elements fits into a `usize`.
80        Some((self.count + len, a))
81    }
82
83    #[inline]
84    fn nth_back(&mut self, n: usize) -> Option<(IdxType, <I as Iterator>::Item)> {
85        let a = self.iter.nth_back(n)?;
86        let len = IdxType::from_usize(self.iter.len()).unwrap();
87        // Can safely add, `ExactSizeIterator` promises that the number of
88        // elements fits into a `usize`.
89        Some((self.count + len, a))
90    }
91}
92
93impl<I, IdxType> ExactSizeIterator for EnumerateIdx<I, IdxType>
94where
95    I: ExactSizeIterator,
96    IdxType: std::ops::Add<Output = IdxType> + FromPrimitive + std::ops::AddAssign + One + Copy,
97{
98    fn len(&self) -> usize {
99        self.iter.len()
100    }
101}
polars_utils/itertools/enumerate_idx.rs

polars_utils/itertools/
enumerate_idx.rs
