1use super::plumbing::*;
2use super::*;
3use std::cmp;
4use std::iter::Fuse;
5
6/// `Interleave` is an iterator that interleaves elements of iterators
7/// `i` and `j` in one continuous iterator. This struct is created by
8/// the [`interleave()`] method on [`IndexedParallelIterator`]
9///
10/// [`interleave()`]: trait.IndexedParallelIterator.html#method.interleave
11/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
12#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
13#[derive(Debug, Clone)]
14pub struct Interleave<I, J>
15where
16 I: IndexedParallelIterator,
17 J: IndexedParallelIterator<Item = I::Item>,
18{
19 i: I,
20 j: J,
21}
22
23impl<I, J> Interleave<I, J>
24where
25 I: IndexedParallelIterator,
26 J: IndexedParallelIterator<Item = I::Item>,
27{
28 /// Creates a new `Interleave` iterator
29 pub(super) fn new(i: I, j: J) -> Self {
30 Interleave { i, j }
31 }
32}
33
34impl<I, J> ParallelIterator for Interleave<I, J>
35where
36 I: IndexedParallelIterator,
37 J: IndexedParallelIterator<Item = I::Item>,
38{
39 type Item = I::Item;
40
41 fn drive_unindexed<C>(self, consumer: C) -> C::Result
42 where
43 C: Consumer<I::Item>,
44 {
45 bridge(self, consumer)
46 }
47
48 fn opt_len(&self) -> Option<usize> {
49 Some(self.len())
50 }
51}
52
53impl<I, J> IndexedParallelIterator for Interleave<I, J>
54where
55 I: IndexedParallelIterator,
56 J: IndexedParallelIterator<Item = I::Item>,
57{
58 fn drive<C>(self, consumer: C) -> C::Result
59 where
60 C: Consumer<Self::Item>,
61 {
62 bridge(self, consumer)
63 }
64
65 fn len(&self) -> usize {
66 self.i.len().checked_add(self.j.len()).expect("overflow")
67 }
68
69 fn with_producer<CB>(self, callback: CB) -> CB::Output
70 where
71 CB: ProducerCallback<Self::Item>,
72 {
73 let (i_len, j_len) = (self.i.len(), self.j.len());
74 return self.i.with_producer(CallbackI {
75 callback,
76 i_len,
77 j_len,
78 i_next: false,
79 j: self.j,
80 });
81
82 struct CallbackI<CB, J> {
83 callback: CB,
84 i_len: usize,
85 j_len: usize,
86 i_next: bool,
87 j: J,
88 }
89
90 impl<CB, J> ProducerCallback<J::Item> for CallbackI<CB, J>
91 where
92 J: IndexedParallelIterator,
93 CB: ProducerCallback<J::Item>,
94 {
95 type Output = CB::Output;
96
97 fn callback<I>(self, i_producer: I) -> Self::Output
98 where
99 I: Producer<Item = J::Item>,
100 {
101 self.j.with_producer(CallbackJ {
102 i_producer,
103 i_len: self.i_len,
104 j_len: self.j_len,
105 i_next: self.i_next,
106 callback: self.callback,
107 })
108 }
109 }
110
111 struct CallbackJ<CB, I> {
112 callback: CB,
113 i_len: usize,
114 j_len: usize,
115 i_next: bool,
116 i_producer: I,
117 }
118
119 impl<CB, I> ProducerCallback<I::Item> for CallbackJ<CB, I>
120 where
121 I: Producer,
122 CB: ProducerCallback<I::Item>,
123 {
124 type Output = CB::Output;
125
126 fn callback<J>(self, j_producer: J) -> Self::Output
127 where
128 J: Producer<Item = I::Item>,
129 {
130 let producer = InterleaveProducer::new(
131 self.i_producer,
132 j_producer,
133 self.i_len,
134 self.j_len,
135 self.i_next,
136 );
137 self.callback.callback(producer)
138 }
139 }
140 }
141}
142
143struct InterleaveProducer<I, J>
144where
145 I: Producer,
146 J: Producer<Item = I::Item>,
147{
148 i: I,
149 j: J,
150 i_len: usize,
151 j_len: usize,
152 i_next: bool,
153}
154
155impl<I, J> InterleaveProducer<I, J>
156where
157 I: Producer,
158 J: Producer<Item = I::Item>,
159{
160 fn new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J> {
161 InterleaveProducer {
162 i,
163 j,
164 i_len,
165 j_len,
166 i_next,
167 }
168 }
169}
170
171impl<I, J> Producer for InterleaveProducer<I, J>
172where
173 I: Producer,
174 J: Producer<Item = I::Item>,
175{
176 type Item = I::Item;
177 type IntoIter = InterleaveSeq<I::IntoIter, J::IntoIter>;
178
179 fn into_iter(self) -> Self::IntoIter {
180 InterleaveSeq {
181 i: self.i.into_iter().fuse(),
182 j: self.j.into_iter().fuse(),
183 i_next: self.i_next,
184 }
185 }
186
187 fn min_len(&self) -> usize {
188 cmp::max(self.i.min_len(), self.j.min_len())
189 }
190
191 fn max_len(&self) -> usize {
192 cmp::min(self.i.max_len(), self.j.max_len())
193 }
194
195 /// We know 0 < index <= self.i_len + self.j_len
196 ///
197 /// Find a, b satisfying:
198 ///
199 /// (1) 0 < a <= self.i_len
200 /// (2) 0 < b <= self.j_len
201 /// (3) a + b == index
202 ///
203 /// For even splits, set a = b = index/2.
204 /// For odd splits, set a = (index/2)+1, b = index/2, if `i`
205 /// should yield the next element, otherwise, if `j` should yield
206 /// the next element, set a = index/2 and b = (index/2)+1
207 fn split_at(self, index: usize) -> (Self, Self) {
208 #[inline]
209 fn odd_offset(flag: bool) -> usize {
210 (!flag) as usize
211 }
212
213 let even = index % 2 == 0;
214 let idx = index >> 1;
215
216 // desired split
217 let (i_idx, j_idx) = (
218 idx + odd_offset(even || self.i_next),
219 idx + odd_offset(even || !self.i_next),
220 );
221
222 let (i_split, j_split) = if self.i_len >= i_idx && self.j_len >= j_idx {
223 (i_idx, j_idx)
224 } else if self.i_len >= i_idx {
225 // j too short
226 (index - self.j_len, self.j_len)
227 } else {
228 // i too short
229 (self.i_len, index - self.i_len)
230 };
231
232 let trailing_i_next = even == self.i_next;
233 let (i_left, i_right) = self.i.split_at(i_split);
234 let (j_left, j_right) = self.j.split_at(j_split);
235
236 (
237 InterleaveProducer::new(i_left, j_left, i_split, j_split, self.i_next),
238 InterleaveProducer::new(
239 i_right,
240 j_right,
241 self.i_len - i_split,
242 self.j_len - j_split,
243 trailing_i_next,
244 ),
245 )
246 }
247}
248
249/// Wrapper for Interleave to implement DoubleEndedIterator and
250/// ExactSizeIterator.
251///
252/// This iterator is fused.
253struct InterleaveSeq<I, J> {
254 i: Fuse<I>,
255 j: Fuse<J>,
256
257 /// Flag to control which iterator should provide the next element. When
258 /// `false` then `i` produces the next element, otherwise `j` produces the
259 /// next element.
260 i_next: bool,
261}
262
263/// Iterator implementation for InterleaveSeq. This implementation is
264/// taken more or less verbatim from itertools. It is replicated here
265/// (instead of calling itertools directly), because we also need to
266/// implement `DoubledEndedIterator` and `ExactSizeIterator`.
267impl<I, J> Iterator for InterleaveSeq<I, J>
268where
269 I: Iterator,
270 J: Iterator<Item = I::Item>,
271{
272 type Item = I::Item;
273
274 #[inline]
275 fn next(&mut self) -> Option<Self::Item> {
276 self.i_next = !self.i_next;
277 if self.i_next {
278 match self.i.next() {
279 None => self.j.next(),
280 r => r,
281 }
282 } else {
283 match self.j.next() {
284 None => self.i.next(),
285 r => r,
286 }
287 }
288 }
289
290 fn size_hint(&self) -> (usize, Option<usize>) {
291 let (ih, jh) = (self.i.size_hint(), self.j.size_hint());
292 let min = ih.0.saturating_add(jh.0);
293 let max = match (ih.1, jh.1) {
294 (Some(x), Some(y)) => x.checked_add(y),
295 _ => None,
296 };
297 (min, max)
298 }
299}
300
301// The implementation for DoubleEndedIterator requires
302// ExactSizeIterator to provide `next_back()`. The last element will
303// come from the iterator that runs out last (ie has the most elements
304// in it). If the iterators have the same number of elements, then the
305// last iterator will provide the last element.
306impl<I, J> DoubleEndedIterator for InterleaveSeq<I, J>
307where
308 I: DoubleEndedIterator + ExactSizeIterator,
309 J: DoubleEndedIterator<Item = I::Item> + ExactSizeIterator<Item = I::Item>,
310{
311 #[inline]
312 fn next_back(&mut self) -> Option<I::Item> {
313 match self.i.len().cmp(&self.j.len()) {
314 Ordering::Less => self.j.next_back(),
315 Ordering::Equal => {
316 if self.i_next {
317 self.i.next_back()
318 } else {
319 self.j.next_back()
320 }
321 }
322 Ordering::Greater => self.i.next_back(),
323 }
324 }
325}
326
327impl<I, J> ExactSizeIterator for InterleaveSeq<I, J>
328where
329 I: ExactSizeIterator,
330 J: ExactSizeIterator<Item = I::Item>,
331{
332 #[inline]
333 fn len(&self) -> usize {
334 self.i.len() + self.j.len()
335 }
336}
337