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