summaryrefslogtreecommitdiffstats
path: root/third_party/rust/rayon/src/iter/enumerate.rs
blob: 980ee7cc7d0d277087dc3cc24ebfbe60e214d5c6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
use super::plumbing::*;
use super::*;
use std::iter;
use std::ops::Range;
use std::usize;

/// `Enumerate` is an iterator that returns the current count along with the element.
/// This struct is created by the [`enumerate()`] method on [`IndexedParallelIterator`]
///
/// [`enumerate()`]: trait.IndexedParallelIterator.html#method.enumerate
/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
#[derive(Debug, Clone)]
pub struct Enumerate<I: IndexedParallelIterator> {
    base: I,
}

impl<I> Enumerate<I>
where
    I: IndexedParallelIterator,
{
    /// Creates a new `Enumerate` iterator.
    pub(super) fn new(base: I) -> Self {
        Enumerate { base }
    }
}

impl<I> ParallelIterator for Enumerate<I>
where
    I: IndexedParallelIterator,
{
    type Item = (usize, I::Item);

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
    where
        C: UnindexedConsumer<Self::Item>,
    {
        bridge(self, consumer)
    }

    fn opt_len(&self) -> Option<usize> {
        Some(self.len())
    }
}

impl<I> IndexedParallelIterator for Enumerate<I>
where
    I: IndexedParallelIterator,
{
    fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
        bridge(self, consumer)
    }

    fn len(&self) -> usize {
        self.base.len()
    }

    fn with_producer<CB>(self, callback: CB) -> CB::Output
    where
        CB: ProducerCallback<Self::Item>,
    {
        return self.base.with_producer(Callback { callback });

        struct Callback<CB> {
            callback: CB,
        }

        impl<I, CB> ProducerCallback<I> for Callback<CB>
        where
            CB: ProducerCallback<(usize, I)>,
        {
            type Output = CB::Output;
            fn callback<P>(self, base: P) -> CB::Output
            where
                P: Producer<Item = I>,
            {
                let producer = EnumerateProducer { base, offset: 0 };
                self.callback.callback(producer)
            }
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// Producer implementation

struct EnumerateProducer<P> {
    base: P,
    offset: usize,
}

impl<P> Producer for EnumerateProducer<P>
where
    P: Producer,
{
    type Item = (usize, P::Item);
    type IntoIter = iter::Zip<Range<usize>, P::IntoIter>;

    fn into_iter(self) -> Self::IntoIter {
        // Enumerate only works for IndexedParallelIterators. Since those
        // have a max length of usize::MAX, their max index is
        // usize::MAX - 1, so the range 0..usize::MAX includes all
        // possible indices.
        //
        // However, we should to use a precise end to the range, otherwise
        // reversing the iterator may have to walk back a long ways before
        // `Zip::next_back` can produce anything.
        let base = self.base.into_iter();
        let end = self.offset + base.len();
        (self.offset..end).zip(base)
    }

    fn min_len(&self) -> usize {
        self.base.min_len()
    }
    fn max_len(&self) -> usize {
        self.base.max_len()
    }

    fn split_at(self, index: usize) -> (Self, Self) {
        let (left, right) = self.base.split_at(index);
        (
            EnumerateProducer {
                base: left,
                offset: self.offset,
            },
            EnumerateProducer {
                base: right,
                offset: self.offset + index,
            },
        )
    }
}