1 use std::cmp::min;
2 
3 use super::plumbing::*;
4 use super::*;
5 use crate::math::div_round_up;
6 
7 /// `Chunks` is an iterator that groups elements of an underlying iterator.
8 ///
9 /// This struct is created by the [`chunks()`] method on [`IndexedParallelIterator`]
10 ///
11 /// [`chunks()`]: trait.IndexedParallelIterator.html#method.chunks
12 /// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
13 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
14 #[derive(Debug, Clone)]
15 pub struct Chunks<I>
16 where
17     I: IndexedParallelIterator,
18 {
19     size: usize,
20     i: I,
21 }
22 
23 impl<I> Chunks<I>
24 where
25     I: IndexedParallelIterator,
26 {
27     /// Creates a new `Chunks` iterator
new(i: I, size: usize) -> Self28     pub(super) fn new(i: I, size: usize) -> Self {
29         Chunks { i, size }
30     }
31 }
32 
33 impl<I> ParallelIterator for Chunks<I>
34 where
35     I: IndexedParallelIterator,
36 {
37     type Item = Vec<I::Item>;
38 
drive_unindexed<C>(self, consumer: C) -> C::Result where C: Consumer<Vec<I::Item>>,39     fn drive_unindexed<C>(self, consumer: C) -> C::Result
40     where
41         C: Consumer<Vec<I::Item>>,
42     {
43         bridge(self, consumer)
44     }
45 
opt_len(&self) -> Option<usize>46     fn opt_len(&self) -> Option<usize> {
47         Some(self.len())
48     }
49 }
50 
51 impl<I> IndexedParallelIterator for Chunks<I>
52 where
53     I: IndexedParallelIterator,
54 {
drive<C>(self, consumer: C) -> C::Result where C: Consumer<Self::Item>,55     fn drive<C>(self, consumer: C) -> C::Result
56     where
57         C: Consumer<Self::Item>,
58     {
59         bridge(self, consumer)
60     }
61 
len(&self) -> usize62     fn len(&self) -> usize {
63         div_round_up(self.i.len(), self.size)
64     }
65 
with_producer<CB>(self, callback: CB) -> CB::Output where CB: ProducerCallback<Self::Item>,66     fn with_producer<CB>(self, callback: CB) -> CB::Output
67     where
68         CB: ProducerCallback<Self::Item>,
69     {
70         let len = self.i.len();
71         return self.i.with_producer(Callback {
72             size: self.size,
73             len,
74             callback,
75         });
76 
77         struct Callback<CB> {
78             size: usize,
79             len: usize,
80             callback: CB,
81         }
82 
83         impl<T, CB> ProducerCallback<T> for Callback<CB>
84         where
85             CB: ProducerCallback<Vec<T>>,
86         {
87             type Output = CB::Output;
88 
89             fn callback<P>(self, base: P) -> CB::Output
90             where
91                 P: Producer<Item = T>,
92             {
93                 self.callback.callback(ChunkProducer {
94                     chunk_size: self.size,
95                     len: self.len,
96                     base,
97                 })
98             }
99         }
100     }
101 }
102 
103 struct ChunkProducer<P>
104 where
105     P: Producer,
106 {
107     chunk_size: usize,
108     len: usize,
109     base: P,
110 }
111 
112 impl<P> Producer for ChunkProducer<P>
113 where
114     P: Producer,
115 {
116     type Item = Vec<P::Item>;
117     type IntoIter = ChunkSeq<P>;
118 
into_iter(self) -> Self::IntoIter119     fn into_iter(self) -> Self::IntoIter {
120         ChunkSeq {
121             chunk_size: self.chunk_size,
122             len: self.len,
123             inner: if self.len > 0 { Some(self.base) } else { None },
124         }
125     }
126 
split_at(self, index: usize) -> (Self, Self)127     fn split_at(self, index: usize) -> (Self, Self) {
128         let elem_index = min(index * self.chunk_size, self.len);
129         let (left, right) = self.base.split_at(elem_index);
130         (
131             ChunkProducer {
132                 chunk_size: self.chunk_size,
133                 len: elem_index,
134                 base: left,
135             },
136             ChunkProducer {
137                 chunk_size: self.chunk_size,
138                 len: self.len - elem_index,
139                 base: right,
140             },
141         )
142     }
143 
min_len(&self) -> usize144     fn min_len(&self) -> usize {
145         div_round_up(self.base.min_len(), self.chunk_size)
146     }
147 
max_len(&self) -> usize148     fn max_len(&self) -> usize {
149         self.base.max_len() / self.chunk_size
150     }
151 }
152 
153 struct ChunkSeq<P> {
154     chunk_size: usize,
155     len: usize,
156     inner: Option<P>,
157 }
158 
159 impl<P> Iterator for ChunkSeq<P>
160 where
161     P: Producer,
162 {
163     type Item = Vec<P::Item>;
164 
next(&mut self) -> Option<Self::Item>165     fn next(&mut self) -> Option<Self::Item> {
166         let producer = self.inner.take()?;
167         if self.len > self.chunk_size {
168             let (left, right) = producer.split_at(self.chunk_size);
169             self.inner = Some(right);
170             self.len -= self.chunk_size;
171             Some(left.into_iter().collect())
172         } else {
173             debug_assert!(self.len > 0);
174             self.len = 0;
175             Some(producer.into_iter().collect())
176         }
177     }
178 
size_hint(&self) -> (usize, Option<usize>)179     fn size_hint(&self) -> (usize, Option<usize>) {
180         let len = self.len();
181         (len, Some(len))
182     }
183 }
184 
185 impl<P> ExactSizeIterator for ChunkSeq<P>
186 where
187     P: Producer,
188 {
189     #[inline]
len(&self) -> usize190     fn len(&self) -> usize {
191         div_round_up(self.len, self.chunk_size)
192     }
193 }
194 
195 impl<P> DoubleEndedIterator for ChunkSeq<P>
196 where
197     P: Producer,
198 {
next_back(&mut self) -> Option<Self::Item>199     fn next_back(&mut self) -> Option<Self::Item> {
200         let producer = self.inner.take()?;
201         if self.len > self.chunk_size {
202             let mut size = self.len % self.chunk_size;
203             if size == 0 {
204                 size = self.chunk_size;
205             }
206             let (left, right) = producer.split_at(self.len - size);
207             self.inner = Some(left);
208             self.len -= size;
209             Some(right.into_iter().collect())
210         } else {
211             debug_assert!(self.len > 0);
212             self.len = 0;
213             Some(producer.into_iter().collect())
214         }
215     }
216 }
217