1 | use super::plumbing::*; |
2 | use super::*; |
3 | use crate::math::div_round_up; |
4 | use std::iter; |
5 | use 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)] |
14 | pub struct StepBy<I: IndexedParallelIterator> { |
15 | base: I, |
16 | step: usize, |
17 | } |
18 | |
19 | impl<I> StepBy<I> |
20 | where |
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 | |
29 | impl<I> ParallelIterator for StepBy<I> |
30 | where |
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 | |
47 | impl<I> IndexedParallelIterator for StepBy<I> |
48 | where |
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 | |
99 | struct StepByProducer<P> { |
100 | base: P, |
101 | step: usize, |
102 | len: usize, |
103 | } |
104 | |
105 | impl<P> Producer for StepByProducer<P> |
106 | where |
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 | } |
142 | |