Skip to main content

asparit/iter/
step_by.rs

1use std::cmp::min;
2
3use crate::{
4    Consumer, Executor, ExecutorCallback, IndexedParallelIterator, IndexedProducer,
5    IndexedProducerCallback, ParallelIterator, Producer, Reducer, Setup, WithIndexedProducer,
6    WithSetup,
7};
8
9/* StepBy */
10
11pub 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
81/* StepByCallback */
82
83struct 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
105/* StepByProducer */
106
107struct 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}