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/// `StepBy` is an iterator that skips `n` elements between each yield, where `n` is the given step.
11/// This struct is created by the [`step_by()`] method on [`IndexedParallelIterator`]
12///
13/// [`step_by()`]: trait.IndexedParallelIterator.html#method.step_by
14/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
15#[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    /// Creates a new `StepBy` iterator.
27    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
99/// ////////////////////////////////////////////////////////////////////////
100/// Producer implementation
101
102struct 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}