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

rayon/iter/
interleave.rs
