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