1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use std::cmp::min;

use super::plumbing::*;
use super::*;
use crate::math::div_round_up;
use std::iter;
use std::usize;

/// `StepBy` is an iterator that skips `n` elements between each yield, where `n` is the given step.
/// This struct is created by the [`step_by()`] method on [`IndexedParallelIterator`]
///
/// [`step_by()`]: trait.IndexedParallelIterator.html#method.step_by
/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
#[derive(Debug, Clone)]
pub struct StepBy<I: IndexedParallelIterator> {
    base: I,
    step: usize,
}

impl<I> StepBy<I>
where
    I: IndexedParallelIterator,
{
    /// Creates a new `StepBy` iterator.
    pub(super) fn new(base: I, step: usize) -> Self {
        StepBy { base, step }
    }
}

impl<I> ParallelIterator for StepBy<I>
where
    I: IndexedParallelIterator,
{
    type Item = I::Item;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
    where
        C: UnindexedConsumer<Self::Item>,
    {
        bridge(self, consumer)
    }

    fn opt_len(&self) -> Option<usize> {
        Some(self.len())
    }
}

impl<I> IndexedParallelIterator for StepBy<I>
where
    I: IndexedParallelIterator,
{
    fn drive<C: Consumer<Self::Item>>(self, consumer: C) -> C::Result {
        bridge(self, consumer)
    }

    fn len(&self) -> usize {
        div_round_up(self.base.len(), self.step)
    }

    fn with_producer<CB>(self, callback: CB) -> CB::Output
    where
        CB: ProducerCallback<Self::Item>,
    {
        let len = self.base.len();
        return self.base.with_producer(Callback {
            callback,
            step: self.step,
            len,
        });

        struct Callback<CB> {
            callback: CB,
            step: usize,
            len: usize,
        }

        impl<T, CB> ProducerCallback<T> for Callback<CB>
        where
            CB: ProducerCallback<T>,
        {
            type Output = CB::Output;
            fn callback<P>(self, base: P) -> CB::Output
            where
                P: Producer<Item = T>,
            {
                let producer = StepByProducer {
                    base,
                    step: self.step,
                    len: self.len,
                };
                self.callback.callback(producer)
            }
        }
    }
}

/// ////////////////////////////////////////////////////////////////////////
/// Producer implementation

struct StepByProducer<P> {
    base: P,
    step: usize,
    len: usize,
}

impl<P> Producer for StepByProducer<P>
where
    P: Producer,
{
    type Item = P::Item;
    type IntoIter = iter::StepBy<P::IntoIter>;

    fn into_iter(self) -> Self::IntoIter {
        self.base.into_iter().step_by(self.step)
    }

    fn split_at(self, index: usize) -> (Self, Self) {
        let elem_index = min(index * self.step, self.len);

        let (left, right) = self.base.split_at(elem_index);
        (
            StepByProducer {
                base: left,
                step: self.step,
                len: elem_index,
            },
            StepByProducer {
                base: right,
                step: self.step,
                len: self.len - elem_index,
            },
        )
    }

    fn min_len(&self) -> usize {
        div_round_up(self.base.min_len(), self.step)
    }

    fn max_len(&self) -> usize {
        self.base.max_len() / self.step
    }
}