1 use super::plumbing::*;
2 use super::*;
3 use std::cmp;
4 use 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)]
14 pub struct Interleave<I, J>
15 where
16     I: IndexedParallelIterator,
17     J: IndexedParallelIterator<Item = I::Item>,
18 {
19     i: I,
20     j: J,
21 }
22 
23 impl<I, J> Interleave<I, J>
24 where
25     I: IndexedParallelIterator,
26     J: IndexedParallelIterator<Item = I::Item>,
27 {
28     /// Creates a new `Interleave` iterator
new(i: I, j: J) -> Self29     pub(super) fn new(i: I, j: J) -> Self {
30         Interleave { i, j }
31     }
32 }
33 
34 impl<I, J> ParallelIterator for Interleave<I, J>
35 where
36     I: IndexedParallelIterator,
37     J: IndexedParallelIterator<Item = I::Item>,
38 {
39     type Item = I::Item;
40 
drive_unindexed<C>(self, consumer: C) -> C::Result where C: Consumer<I::Item>,41     fn drive_unindexed<C>(self, consumer: C) -> C::Result
42     where
43         C: Consumer<I::Item>,
44     {
45         bridge(self, consumer)
46     }
47 
opt_len(&self) -> Option<usize>48     fn opt_len(&self) -> Option<usize> {
49         Some(self.len())
50     }
51 }
52 
53 impl<I, J> IndexedParallelIterator for Interleave<I, J>
54 where
55     I: IndexedParallelIterator,
56     J: IndexedParallelIterator<Item = I::Item>,
57 {
drive<C>(self, consumer: C) -> C::Result where C: Consumer<Self::Item>,58     fn drive<C>(self, consumer: C) -> C::Result
59     where
60         C: Consumer<Self::Item>,
61     {
62         bridge(self, consumer)
63     }
64 
len(&self) -> usize65     fn len(&self) -> usize {
66         self.i.len().checked_add(self.j.len()).expect("overflow")
67     }
68 
with_producer<CB>(self, callback: CB) -> CB::Output where CB: ProducerCallback<Self::Item>,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 
143 struct InterleaveProducer<I, J>
144 where
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 
155 impl<I, J> InterleaveProducer<I, J>
156 where
157     I: Producer,
158     J: Producer<Item = I::Item>,
159 {
new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J>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 
171 impl<I, J> Producer for InterleaveProducer<I, J>
172 where
173     I: Producer,
174     J: Producer<Item = I::Item>,
175 {
176     type Item = I::Item;
177     type IntoIter = InterleaveSeq<I::IntoIter, J::IntoIter>;
178 
into_iter(self) -> Self::IntoIter179     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 
min_len(&self) -> usize187     fn min_len(&self) -> usize {
188         cmp::max(self.i.min_len(), self.j.min_len())
189     }
190 
max_len(&self) -> usize191     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
split_at(self, index: usize) -> (Self, Self)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.
253 struct 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`.
267 impl<I, J> Iterator for InterleaveSeq<I, J>
268 where
269     I: Iterator,
270     J: Iterator<Item = I::Item>,
271 {
272     type Item = I::Item;
273 
274     #[inline]
next(&mut self) -> Option<Self::Item>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 
size_hint(&self) -> (usize, Option<usize>)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.
306 impl<I, J> DoubleEndedIterator for InterleaveSeq<I, J>
307 where
308     I: DoubleEndedIterator + ExactSizeIterator,
309     J: DoubleEndedIterator<Item = I::Item> + ExactSizeIterator<Item = I::Item>,
310 {
311     #[inline]
next_back(&mut self) -> Option<I::Item>312     fn next_back(&mut self) -> Option<I::Item> {
313         if self.i.len() == self.j.len() {
314             if self.i_next {
315                 self.i.next_back()
316             } else {
317                 self.j.next_back()
318             }
319         } else if self.i.len() < self.j.len() {
320             self.j.next_back()
321         } else {
322             self.i.next_back()
323         }
324     }
325 }
326 
327 impl<I, J> ExactSizeIterator for InterleaveSeq<I, J>
328 where
329     I: ExactSizeIterator,
330     J: ExactSizeIterator<Item = I::Item>,
331 {
332     #[inline]
len(&self) -> usize333     fn len(&self) -> usize {
334         self.i.len() + self.j.len()
335     }
336 }
337