/rust/registry/src/index.crates.io-1949cf8c6b5b557f/rayon-1.11.0/src/iter/interleave.rs

Source
use super::plumbing::*;
use super::*;
use std::iter::Fuse;

/// `Interleave` is an iterator that interleaves elements of iterators
/// `i` and `j` in one continuous iterator. This struct is created by
/// the [`interleave()`] method on [`IndexedParallelIterator`]
///
/// [`interleave()`]: IndexedParallelIterator::interleave()
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
#[derive(Debug, Clone)]
pub struct Interleave<I, J> {
    i: I,
    j: J,
}

impl<I, J> Interleave<I, J> {
    /// Creates a new `Interleave` iterator
    pub(super) fn new(i: I, j: J) -> Self {
        Interleave { i, j }
    }
}

impl<I, J> ParallelIterator for Interleave<I, J>
where
    I: IndexedParallelIterator,
    J: IndexedParallelIterator<Item = I::Item>,
{
    type Item = I::Item;

    fn drive_unindexed<C>(self, consumer: C) -> C::Result
    where
        C: Consumer<I::Item>,
    {
        bridge(self, consumer)
    }

    fn opt_len(&self) -> Option<usize> {
        Some(self.len())
    }
}

impl<I, J> IndexedParallelIterator for Interleave<I, J>
where
    I: IndexedParallelIterator,
    J: IndexedParallelIterator<Item = I::Item>,
{
    fn drive<C>(self, consumer: C) -> C::Result
    where
        C: Consumer<Self::Item>,
    {
        bridge(self, consumer)
    }

    fn len(&self) -> usize {
        self.i.len().checked_add(self.j.len()).expect("overflow")
    }

    fn with_producer<CB>(self, callback: CB) -> CB::Output
    where
        CB: ProducerCallback<Self::Item>,
    {
        let (i_len, j_len) = (self.i.len(), self.j.len());
        return self.i.with_producer(CallbackI {
            callback,
            i_len,
            j_len,
            i_next: false,
            j: self.j,
        });

        struct CallbackI<CB, J> {
            callback: CB,
            i_len: usize,
            j_len: usize,
            i_next: bool,
            j: J,
        }

        impl<CB, J> ProducerCallback<J::Item> for CallbackI<CB, J>
        where
            J: IndexedParallelIterator,
            CB: ProducerCallback<J::Item>,
        {
            type Output = CB::Output;

            fn callback<I>(self, i_producer: I) -> Self::Output
            where
                I: Producer<Item = J::Item>,
            {
                self.j.with_producer(CallbackJ {
                    i_producer,
                    i_len: self.i_len,
                    j_len: self.j_len,
                    i_next: self.i_next,
                    callback: self.callback,
                })
            }
        }

        struct CallbackJ<CB, I> {
            callback: CB,
            i_len: usize,
            j_len: usize,
            i_next: bool,
            i_producer: I,
        }

        impl<CB, I> ProducerCallback<I::Item> for CallbackJ<CB, I>
        where
            I: Producer,
            CB: ProducerCallback<I::Item>,
        {
            type Output = CB::Output;

            fn callback<J>(self, j_producer: J) -> Self::Output
            where
                J: Producer<Item = I::Item>,
            {
                let producer = InterleaveProducer::new(
                    self.i_producer,
                    j_producer,
                    self.i_len,
                    self.j_len,
                    self.i_next,
                );
                self.callback.callback(producer)
            }
        }
    }
}

struct InterleaveProducer<I, J>
where
    I: Producer,
    J: Producer<Item = I::Item>,
{
    i: I,
    j: J,
    i_len: usize,
    j_len: usize,
    i_next: bool,
}

impl<I, J> InterleaveProducer<I, J>
where
    I: Producer,
    J: Producer<Item = I::Item>,
{
    fn new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J> {
        InterleaveProducer {
            i,
            j,
            i_len,
            j_len,
            i_next,
        }
    }
}

impl<I, J> Producer for InterleaveProducer<I, J>
where
    I: Producer,
    J: Producer<Item = I::Item>,
{
    type Item = I::Item;
    type IntoIter = InterleaveSeq<I::IntoIter, J::IntoIter>;

    fn into_iter(self) -> Self::IntoIter {
        InterleaveSeq {
            i: self.i.into_iter().fuse(),
            j: self.j.into_iter().fuse(),
            i_next: self.i_next,
        }
    }

    fn min_len(&self) -> usize {
        Ord::max(self.i.min_len(), self.j.min_len())
    }

    fn max_len(&self) -> usize {
        Ord::min(self.i.max_len(), self.j.max_len())
    }

    /// We know 0 < index <= self.i_len + self.j_len
    ///
    /// Find a, b satisfying:
    ///
    ///  (1) 0 < a <= self.i_len
    ///  (2) 0 < b <= self.j_len
    ///  (3) a + b == index
    ///
    /// For even splits, set a = b = index/2.
    /// For odd splits, set a = (index/2)+1, b = index/2, if `i`
    /// should yield the next element, otherwise, if `j` should yield
    /// the next element, set a = index/2 and b = (index/2)+1
    fn split_at(self, index: usize) -> (Self, Self) {
        #[inline]
        fn odd_offset(flag: bool) -> usize {
            (!flag) as usize
        }

        let even = index % 2 == 0;
        let idx = index >> 1;

        // desired split
        let (i_idx, j_idx) = (
            idx + odd_offset(even || self.i_next),
            idx + odd_offset(even || !self.i_next),
        );

        let (i_split, j_split) = if self.i_len >= i_idx && self.j_len >= j_idx {
            (i_idx, j_idx)
        } else if self.i_len >= i_idx {
            // j too short
            (index - self.j_len, self.j_len)
        } else {
            // i too short
            (self.i_len, index - self.i_len)
        };

        let trailing_i_next = even == self.i_next;
        let (i_left, i_right) = self.i.split_at(i_split);
        let (j_left, j_right) = self.j.split_at(j_split);

        (
            InterleaveProducer::new(i_left, j_left, i_split, j_split, self.i_next),
            InterleaveProducer::new(
                i_right,
                j_right,
                self.i_len - i_split,
                self.j_len - j_split,
                trailing_i_next,
            ),
        )
    }
}

/// Wrapper for Interleave to implement DoubleEndedIterator and
/// ExactSizeIterator.
///
/// This iterator is fused.
struct InterleaveSeq<I, J> {
    i: Fuse<I>,
    j: Fuse<J>,

    /// Flag to control which iterator should provide the next element. When
    /// `false` then `i` produces the next element, otherwise `j` produces the
    /// next element.
    i_next: bool,
}

/// Iterator implementation for InterleaveSeq. This implementation is
/// taken more or less verbatim from itertools. It is replicated here
/// (instead of calling itertools directly), because we also need to
/// implement `DoubledEndedIterator` and `ExactSizeIterator`.
impl<I, J> Iterator for InterleaveSeq<I, J>
where
    I: Iterator,
    J: Iterator<Item = I::Item>,
{
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        self.i_next = !self.i_next;
        if self.i_next {
            match self.i.next() {
                None => self.j.next(),
                r => r,
            }
        } else {
            match self.j.next() {
                None => self.i.next(),
                r => r,
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (ih, jh) = (self.i.size_hint(), self.j.size_hint());
        let min = ih.0.saturating_add(jh.0);
        let max = match (ih.1, jh.1) {
            (Some(x), Some(y)) => x.checked_add(y),
            _ => None,
        };
        (min, max)
    }
}

// The implementation for DoubleEndedIterator requires
// ExactSizeIterator to provide `next_back()`. The last element will
// come from the iterator that runs out last (ie has the most elements
// in it). If the iterators have the same number of elements, then the
// last iterator will provide the last element.
impl<I, J> DoubleEndedIterator for InterleaveSeq<I, J>
where
    I: DoubleEndedIterator + ExactSizeIterator,
    J: DoubleEndedIterator<Item = I::Item> + ExactSizeIterator<Item = I::Item>,
{
    #[inline]
    fn next_back(&mut self) -> Option<I::Item> {
        match self.i.len().cmp(&self.j.len()) {
            Ordering::Less => self.j.next_back(),
            Ordering::Equal => {
                if self.i_next {
                    self.i.next_back()
                } else {
                    self.j.next_back()
                }
            }
            Ordering::Greater => self.i.next_back(),
        }
    }
}

impl<I, J> ExactSizeIterator for InterleaveSeq<I, J>
where
    I: ExactSizeIterator,
    J: ExactSizeIterator<Item = I::Item>,
{
    #[inline]
    fn len(&self) -> usize {
        self.i.len() + self.j.len()
    }
}

Coverage Report

Created: 2025-10-10 07:21

Line	Count	Source
1		use super::plumbing::*;
2		use super::*;
3		use 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)]
12		pub struct Interleave<I, J> {
13		i: I,
14		j: J,
15		}
16
17		impl<I, J> Interleave<I, J> {
18		/// Creates a new `Interleave` iterator
19	0	pub(super) fn new(i: I, j: J) -> Self {
20	0	Interleave { i, j }
21	0	}
22		}
23
24		impl<I, J> ParallelIterator for Interleave<I, J>
25		where
26		I: IndexedParallelIterator,
27		J: IndexedParallelIterator<Item = I::Item>,
28		{
29		type Item = I::Item;
30
31	0	fn drive_unindexed<C>(self, consumer: C) -> C::Result
32	0	where
33	0	C: Consumer<I::Item>,
34		{
35	0	bridge(self, consumer)
36	0	}
37
38	0	fn opt_len(&self) -> Option<usize> {
39	0	Some(self.len())
40	0	}
41		}
42
43		impl<I, J> IndexedParallelIterator for Interleave<I, J>
44		where
45		I: IndexedParallelIterator,
46		J: IndexedParallelIterator<Item = I::Item>,
47		{
48	0	fn drive<C>(self, consumer: C) -> C::Result
49	0	where
50	0	C: Consumer<Self::Item>,
51		{
52	0	bridge(self, consumer)
53	0	}
54
55	0	fn len(&self) -> usize {
56	0	self.i.len().checked_add(self.j.len()).expect("overflow")
57	0	}
58
59	0	fn with_producer<CB>(self, callback: CB) -> CB::Output
60	0	where
61	0	CB: ProducerCallback<Self::Item>,
62		{
63	0	let (i_len, j_len) = (self.i.len(), self.j.len());
64	0	return self.i.with_producer(CallbackI {
65	0	callback,
66	0	i_len,
67	0	j_len,
68	0	i_next: false,
69	0	j: self.j,
70	0	});
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	0	fn callback<I>(self, i_producer: I) -> Self::Output
88	0	where
89	0	I: Producer<Item = J::Item>,
90		{
91	0	self.j.with_producer(CallbackJ {
92	0	i_producer,
93	0	i_len: self.i_len,
94	0	j_len: self.j_len,
95	0	i_next: self.i_next,
96	0	callback: self.callback,
97	0	})
98	0	}
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	0	fn callback<J>(self, j_producer: J) -> Self::Output
117	0	where
118	0	J: Producer<Item = I::Item>,
119		{
120	0	let producer = InterleaveProducer::new(
121	0	self.i_producer,
122	0	j_producer,
123	0	self.i_len,
124	0	self.j_len,
125	0	self.i_next,
126		);
127	0	self.callback.callback(producer)
128	0	}
129		}
130	0	}
131		}
132
133		struct InterleaveProducer<I, J>
134		where
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
145		impl<I, J> InterleaveProducer<I, J>
146		where
147		I: Producer,
148		J: Producer<Item = I::Item>,
149		{
150	0	fn new(i: I, j: J, i_len: usize, j_len: usize, i_next: bool) -> InterleaveProducer<I, J> {
151	0	InterleaveProducer {
152	0	i,
153	0	j,
154	0	i_len,
155	0	j_len,
156	0	i_next,
157	0	}
158	0	}
159		}
160
161		impl<I, J> Producer for InterleaveProducer<I, J>
162		where
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	0	fn into_iter(self) -> Self::IntoIter {
170	0	InterleaveSeq {
171	0	i: self.i.into_iter().fuse(),
172	0	j: self.j.into_iter().fuse(),
173	0	i_next: self.i_next,
174	0	}
175	0	}
176
177	0	fn min_len(&self) -> usize {
178	0	Ord::max(self.i.min_len(), self.j.min_len())
179	0	}
180
181	0	fn max_len(&self) -> usize {
182	0	Ord::min(self.i.max_len(), self.j.max_len())
183	0	}
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	0	fn split_at(self, index: usize) -> (Self, Self) {
198		#[inline]
199	0	fn odd_offset(flag: bool) -> usize {
200	0	(!flag) as usize
201	0	}
202
203	0	let even = index % 2 == 0;
204	0	let idx = index >> 1;
205
206		// desired split
207	0	let (i_idx, j_idx) = (
208	0	idx + odd_offset(even \|\| self.i_next),
209	0	idx + odd_offset(even \|\| !self.i_next),
210		);
211
212	0	let (i_split, j_split) = if self.i_len >= i_idx && self.j_len >= j_idx {
213	0	(i_idx, j_idx)
214	0	} else if self.i_len >= i_idx {
215		// j too short
216	0	(index - self.j_len, self.j_len)
217		} else {
218		// i too short
219	0	(self.i_len, index - self.i_len)
220		};
221
222	0	let trailing_i_next = even == self.i_next;
223	0	let (i_left, i_right) = self.i.split_at(i_split);
224	0	let (j_left, j_right) = self.j.split_at(j_split);
225
226	0	(
227	0	InterleaveProducer::new(i_left, j_left, i_split, j_split, self.i_next),
228	0	InterleaveProducer::new(
229	0	i_right,
230	0	j_right,
231	0	self.i_len - i_split,
232	0	self.j_len - j_split,
233	0	trailing_i_next,
234	0	),
235	0	)
236	0	}
237		}
238
239		/// Wrapper for Interleave to implement DoubleEndedIterator and
240		/// ExactSizeIterator.
241		///
242		/// This iterator is fused.
243		struct 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`.
257		impl<I, J> Iterator for InterleaveSeq<I, J>
258		where
259		I: Iterator,
260		J: Iterator<Item = I::Item>,
261		{
262		type Item = I::Item;
263
264		#[inline]
265	0	fn next(&mut self) -> Option<Self::Item> {
266	0	self.i_next = !self.i_next;
267	0	if self.i_next {
268	0	match self.i.next() {
269	0	None => self.j.next(),
270	0	r => r,
271		}
272		} else {
273	0	match self.j.next() {
274	0	None => self.i.next(),
275	0	r => r,
276		}
277		}
278	0	}
279
280	0	fn size_hint(&self) -> (usize, Option<usize>) {
281	0	let (ih, jh) = (self.i.size_hint(), self.j.size_hint());
282	0	let min = ih.0.saturating_add(jh.0);
283	0	let max = match (ih.1, jh.1) {
284	0	(Some(x), Some(y)) => x.checked_add(y),
285	0	_ => None,
286		};
287	0	(min, max)
288	0	}
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.
296		impl<I, J> DoubleEndedIterator for InterleaveSeq<I, J>
297		where
298		I: DoubleEndedIterator + ExactSizeIterator,
299		J: DoubleEndedIterator<Item = I::Item> + ExactSizeIterator<Item = I::Item>,
300		{
301		#[inline]
302	0	fn next_back(&mut self) -> Option<I::Item> {
303	0	match self.i.len().cmp(&self.j.len()) {
304	0	Ordering::Less => self.j.next_back(),
305		Ordering::Equal => {
306	0	if self.i_next {
307	0	self.i.next_back()
308		} else {
309	0	self.j.next_back()
310		}
311		}
312	0	Ordering::Greater => self.i.next_back(),
313		}
314	0	}
315		}
316
317		impl<I, J> ExactSizeIterator for InterleaveSeq<I, J>
318		where
319		I: ExactSizeIterator,
320		J: ExactSizeIterator<Item = I::Item>,
321		{
322		#[inline]
323	0	fn len(&self) -> usize {
324	0	self.i.len() + self.j.len()
325	0	}
326		}