rayon/iter/
step_by.rs

1use super::plumbing::*;
2use super::*;
3use std::iter;
4
5/// `StepBy` is an iterator that skips `n` elements between each yield, where `n` is the given step.
6/// This struct is created by the [`step_by()`] method on [`IndexedParallelIterator`]
7///
8/// [`step_by()`]: IndexedParallelIterator::step_by()
9#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
10#[derive(Debug, Clone)]
11pub struct StepBy<I> {
12    base: I,
13    step: usize,
14}
15
16impl<I> StepBy<I> {
17    /// Creates a new `StepBy` iterator.
18    pub(super) fn new(base: I, step: usize) -> Self {
19        StepBy { base, step }
20    }
21}
22
23impl<I> ParallelIterator for StepBy<I>
24where
25    I: IndexedParallelIterator,
26{
27    type Item = 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 StepBy<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().div_ceil(self.step)
51    }
52
53    fn with_producer<CB>(self, callback: CB) -> CB::Output
54    where
55        CB: ProducerCallback<Self::Item>,
56    {
57        let len = self.base.len();
58        return self.base.with_producer(Callback {
59            callback,
60            step: self.step,
61            len,
62        });
63
64        struct Callback<CB> {
65            callback: CB,
66            step: usize,
67            len: usize,
68        }
69
70        impl<T, CB> ProducerCallback<T> for Callback<CB>
71        where
72            CB: ProducerCallback<T>,
73        {
74            type Output = CB::Output;
75            fn callback<P>(self, base: P) -> CB::Output
76            where
77                P: Producer<Item = T>,
78            {
79                let producer = StepByProducer {
80                    base,
81                    step: self.step,
82                    len: self.len,
83                };
84                self.callback.callback(producer)
85            }
86        }
87    }
88}
89
90// ////////////////////////////////////////////////////////////////////////
91// Producer implementation
92
93struct StepByProducer<P> {
94    base: P,
95    step: usize,
96    len: usize,
97}
98
99impl<P> Producer for StepByProducer<P>
100where
101    P: Producer,
102{
103    type Item = P::Item;
104    type IntoIter = iter::StepBy<P::IntoIter>;
105
106    fn into_iter(self) -> Self::IntoIter {
107        self.base.into_iter().step_by(self.step)
108    }
109
110    fn split_at(self, index: usize) -> (Self, Self) {
111        let elem_index = Ord::min(index * self.step, self.len);
112
113        let (left, right) = self.base.split_at(elem_index);
114        (
115            StepByProducer {
116                base: left,
117                step: self.step,
118                len: elem_index,
119            },
120            StepByProducer {
121                base: right,
122                step: self.step,
123                len: self.len - elem_index,
124            },
125        )
126    }
127
128    fn min_len(&self) -> usize {
129        self.base.min_len().div_ceil(self.step)
130    }
131
132    fn max_len(&self) -> usize {
133        self.base.max_len() / self.step
134    }
135}
rayon/iter/step_by.rs

rayon/iter/
step_by.rs
