1use std::cmp::min;
2
3use super::plumbing::*;
4use super::*;
5use crate::math::div_round_up;
6use std::iter;
7use 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)]
16pub struct StepBy<I: IndexedParallelIterator> {
17 base: I,
18 step: usize,
19}
20
21impl<I> StepBy<I>
22where
23 I: IndexedParallelIterator,
24{
25 /// Creates a new `StepBy` iterator.
26 pub(super) fn new(base: I, step: usize) -> Self {
27 StepBy { base, step }
28 }
29}
30
31impl<I> ParallelIterator for StepBy<I>
32where
33 I: IndexedParallelIterator,
34{
35 type Item = I::Item;
36
37 fn drive_unindexed<C>(self, consumer: C) -> C::Result
38 where
39 C: UnindexedConsumer<Self::Item>,
40 {
41 bridge(self, consumer)
42 }
43
44 fn opt_len(&self) -> Option<usize> {
45 Some(self.len())
46 }
47}
48
49impl<I> IndexedParallelIterator for StepBy<I>
50where
51 I: IndexedParallelIterator,
52{
53 fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
54 bridge(self, consumer)
55 }
56
57 fn len(&self) -> usize {
58 div_round_up(self.base.len(), self.step)
59 }
60
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
101struct StepByProducer<P> {
102 base: P,
103 step: usize,
104 len: usize,
105}
106
107impl<P> Producer for StepByProducer<P>
108where
109 P: Producer,
110{
111 type Item = P::Item;
112 type IntoIter = iter::StepBy<P::IntoIter>;
113
114 fn into_iter(self) -> Self::IntoIter {
115 self.base.into_iter().step_by(self.step)
116 }
117
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
136 fn min_len(&self) -> usize {
137 div_round_up(self.base.min_len(), self.step)
138 }
139
140 fn max_len(&self) -> usize {
141 self.base.max_len() / self.step
142 }
143}
144