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