rayon/iter/
step_by.rs

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