| 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 | |