rayon/iter/
interleave.rs

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}
rayon/iter/interleave.rs

rayon/iter/
interleave.rs
