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