1use std::cmp::min;
2
3use crate::{
4 Consumer, Executor, ExecutorCallback, IndexedParallelIterator, IndexedProducer,
5 IndexedProducerCallback, ParallelIterator, Producer, Reducer, Setup, WithIndexedProducer,
6 WithSetup,
7};
8
9pub struct StepBy<X> {
12 base: X,
13 step: usize,
14}
15
16impl<X> StepBy<X> {
17 pub fn new(base: X, step: usize) -> Self {
18 Self { base, step }
19 }
20}
21
22impl<'a, X, I> ParallelIterator<'a> for StepBy<X>
23where
24 X: IndexedParallelIterator<'a, Item = I> + WithIndexedProducer<'a, Item = I>,
25 I: Send + 'a,
26{
27 type Item = I;
28
29 fn drive<E, C, D, R>(self, executor: E, consumer: C) -> E::Result
30 where
31 E: Executor<'a, D>,
32 C: Consumer<Self::Item, Result = D, Reducer = R> + 'a,
33 D: Send + 'a,
34 R: Reducer<D> + Send + 'a,
35 {
36 self.with_indexed_producer(ExecutorCallback::new(executor, consumer))
37 }
38
39 fn len_hint_opt(&self) -> Option<usize> {
40 self.base.len_hint_opt()
41 }
42}
43
44impl<'a, X, I> IndexedParallelIterator<'a> for StepBy<X>
45where
46 X: IndexedParallelIterator<'a, Item = I> + WithIndexedProducer<'a, Item = I>,
47 I: Send + 'a,
48{
49 fn drive_indexed<E, C, D, R>(self, executor: E, consumer: C) -> E::Result
50 where
51 E: Executor<'a, D>,
52 C: Consumer<Self::Item, Result = D, Reducer = R> + 'a,
53 D: Send + 'a,
54 R: Reducer<D> + Send + 'a,
55 {
56 self.with_indexed_producer(ExecutorCallback::new(executor, consumer))
57 }
58
59 fn len_hint(&self) -> usize {
60 self.base.len_hint()
61 }
62}
63
64impl<'a, X> WithIndexedProducer<'a> for StepBy<X>
65where
66 X: WithIndexedProducer<'a>,
67{
68 type Item = X::Item;
69
70 fn with_indexed_producer<CB>(self, base: CB) -> CB::Output
71 where
72 CB: IndexedProducerCallback<'a, Self::Item>,
73 {
74 self.base.with_indexed_producer(StepByCallback {
75 base,
76 step: self.step,
77 })
78 }
79}
80
81struct StepByCallback<CB> {
84 base: CB,
85 step: usize,
86}
87
88impl<'a, CB, I> IndexedProducerCallback<'a, I> for StepByCallback<CB>
89where
90 CB: IndexedProducerCallback<'a, I>,
91{
92 type Output = CB::Output;
93
94 fn callback<P>(self, base: P) -> Self::Output
95 where
96 P: IndexedProducer<Item = I> + 'a,
97 {
98 self.base.callback(StepByProducer {
99 base,
100 step: self.step,
101 })
102 }
103}
104
105struct StepByProducer<P> {
108 base: P,
109 step: usize,
110}
111
112impl<P> WithSetup for StepByProducer<P>
113where
114 P: WithSetup,
115{
116 fn setup(&self) -> Setup {
117 let Setup {
118 splits,
119 min_len,
120 max_len,
121 } = self.base.setup();
122
123 Setup {
124 splits,
125 min_len: min_len.map(|x| if x > 0 { (x - 1) / self.step + 1 } else { x }),
126 max_len: max_len.map(|x| x / self.step),
127 }
128 }
129}
130
131impl<P> Producer for StepByProducer<P>
132where
133 P: IndexedProducer,
134{
135 type Item = P::Item;
136 type IntoIter = std::iter::StepBy<P::IntoIter>;
137
138 fn into_iter(self) -> Self::IntoIter {
139 self.base.into_iter().step_by(self.step)
140 }
141
142 fn split(self) -> (Self, Option<Self>) {
143 let len = self.len();
144 if len < 2 {
145 return (self, None);
146 }
147
148 let (left, right) = self.split_at(len / 2);
149
150 (left, Some(right))
151 }
152}
153
154impl<P> IndexedProducer for StepByProducer<P>
155where
156 P: IndexedProducer,
157{
158 type Item = P::Item;
159 type IntoIter = std::iter::StepBy<P::IntoIter>;
160
161 fn into_iter(self) -> Self::IntoIter {
162 self.base.into_iter().step_by(self.step)
163 }
164
165 fn len(&self) -> usize {
166 self.base.len() / self.step
167 }
168
169 fn split_at(self, index: usize) -> (Self, Self) {
170 let index = min(index * self.step, self.base.len());
171
172 let (left, right) = self.base.split_at(index);
173
174 let left = Self {
175 base: left,
176 step: self.step,
177 };
178 let right = Self {
179 base: right,
180 step: self.step,
181 };
182
183 (left, right)
184 }
185}