rayon/iter/
enumerate.rs

1use super::plumbing::*;
2use super::*;
3use std::iter;
4use std::ops::Range;
5
6/// `Enumerate` is an iterator that returns the current count along with the element.
7/// This struct is created by the [`enumerate()`] method on [`IndexedParallelIterator`]
8///
9/// [`enumerate()`]: IndexedParallelIterator::enumerate()
10#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
11#[derive(Debug, Clone)]
12pub struct Enumerate<I> {
13    base: I,
14}
15
16impl<I> Enumerate<I> {
17    /// Creates a new `Enumerate` iterator.
18    pub(super) fn new(base: I) -> Self {
19        Enumerate { base }
20    }
21}
22
23impl<I> ParallelIterator for Enumerate<I>
24where
25    I: IndexedParallelIterator,
26{
27    type Item = (usize, I::Item);
28
29    fn drive_unindexed<C>(self, consumer: C) -> C::Result
30    where
31        C: UnindexedConsumer<Self::Item>,
32    {
33        bridge(self, consumer)
34    }
35
36    fn opt_len(&self) -> Option<usize> {
37        Some(self.len())
38    }
39}
40
41impl<I> IndexedParallelIterator for Enumerate<I>
42where
43    I: IndexedParallelIterator,
44{
45    fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
46        bridge(self, consumer)
47    }
48
49    fn len(&self) -> usize {
50        self.base.len()
51    }
52
53    fn with_producer<CB>(self, callback: CB) -> CB::Output
54    where
55        CB: ProducerCallback<Self::Item>,
56    {
57        return self.base.with_producer(Callback { callback });
58
59        struct Callback<CB> {
60            callback: CB,
61        }
62
63        impl<I, CB> ProducerCallback<I> for Callback<CB>
64        where
65            CB: ProducerCallback<(usize, I)>,
66        {
67            type Output = CB::Output;
68            fn callback<P>(self, base: P) -> CB::Output
69            where
70                P: Producer<Item = I>,
71            {
72                let producer = EnumerateProducer { base, offset: 0 };
73                self.callback.callback(producer)
74            }
75        }
76    }
77}
78
79// ////////////////////////////////////////////////////////////////////////
80// Producer implementation
81
82struct EnumerateProducer<P> {
83    base: P,
84    offset: usize,
85}
86
87impl<P> Producer for EnumerateProducer<P>
88where
89    P: Producer,
90{
91    type Item = (usize, P::Item);
92    type IntoIter = iter::Zip<Range<usize>, P::IntoIter>;
93
94    fn into_iter(self) -> Self::IntoIter {
95        // Enumerate only works for IndexedParallelIterators. Since those
96        // have a max length of usize::MAX, their max index is
97        // usize::MAX - 1, so the range 0..usize::MAX includes all
98        // possible indices.
99        //
100        // However, we should to use a precise end to the range, otherwise
101        // reversing the iterator may have to walk back a long ways before
102        // `Zip::next_back` can produce anything.
103        let base = self.base.into_iter();
104        let end = self.offset + base.len();
105        (self.offset..end).zip(base)
106    }
107
108    fn min_len(&self) -> usize {
109        self.base.min_len()
110    }
111    fn max_len(&self) -> usize {
112        self.base.max_len()
113    }
114
115    fn split_at(self, index: usize) -> (Self, Self) {
116        let (left, right) = self.base.split_at(index);
117        (
118            EnumerateProducer {
119                base: left,
120                offset: self.offset,
121            },
122            EnumerateProducer {
123                base: right,
124                offset: self.offset + index,
125            },
126        )
127    }
128}