1use std::cmp::Ordering;
2use std::fmt;
3use std::iter::{Fuse, FusedIterator};
4use std::marker::PhantomData;
5
6use either::Either;
7
8use super::adaptors::{put_back, PutBack};
9use crate::either_or_both::EitherOrBoth;
10use crate::size_hint::{self, SizeHint};
11#[cfg(doc)]
12use crate::Itertools;
13
14#[derive(Clone, Debug)]
15pub struct MergeLte;
16
17/// An iterator adaptor that merges the two base iterators in ascending order.
18/// If both base iterators are sorted (ascending), the result is sorted.
19///
20/// Iterator element type is `I::Item`.
21///
22/// See [`.merge()`](crate::Itertools::merge_by) for more information.
23pub type Merge<I, J> = MergeBy<I, J, MergeLte>;
24
25/// Create an iterator that merges elements in `i` and `j`.
26///
27/// [`IntoIterator`] enabled version of [`Itertools::merge`](crate::Itertools::merge).
28///
29/// ```
30/// use itertools::merge;
31///
32/// for elt in merge(&[1, 2, 3], &[2, 3, 4]) {
33/// /* loop body */
34/// }
35/// ```
36pub fn merge<I, J>(
37 i: I,
38 j: J,
39) -> Merge<<I as IntoIterator>::IntoIter, <J as IntoIterator>::IntoIter>
40where
41 I: IntoIterator,
42 J: IntoIterator<Item = I::Item>,
43 I::Item: PartialOrd,
44{
45 merge_by_new(a:i, b:j, cmp:MergeLte)
46}
47
48/// An iterator adaptor that merges the two base iterators in ascending order.
49/// If both base iterators are sorted (ascending), the result is sorted.
50///
51/// Iterator element type is `I::Item`.
52///
53/// See [`.merge_by()`](crate::Itertools::merge_by) for more information.
54#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
55pub struct MergeBy<I: Iterator, J: Iterator, F> {
56 left: PutBack<Fuse<I>>,
57 right: PutBack<Fuse<J>>,
58 cmp_fn: F,
59}
60
61/// Create a `MergeBy` iterator.
62pub fn merge_by_new<I, J, F>(a: I, b: J, cmp: F) -> MergeBy<I::IntoIter, J::IntoIter, F>
63where
64 I: IntoIterator,
65 J: IntoIterator<Item = I::Item>,
66{
67 MergeBy {
68 left: put_back(iterable:a.into_iter().fuse()),
69 right: put_back(iterable:b.into_iter().fuse()),
70 cmp_fn: cmp,
71 }
72}
73
74/// Return an iterator adaptor that merge-joins items from the two base iterators in ascending order.
75///
76/// [`IntoIterator`] enabled version of [`Itertools::merge_join_by`].
77pub fn merge_join_by<I, J, F, T>(
78 left: I,
79 right: J,
80 cmp_fn: F,
81) -> MergeJoinBy<I::IntoIter, J::IntoIter, F>
82where
83 I: IntoIterator,
84 J: IntoIterator,
85 F: FnMut(&I::Item, &J::Item) -> T,
86{
87 MergeBy {
88 left: put_back(iterable:left.into_iter().fuse()),
89 right: put_back(iterable:right.into_iter().fuse()),
90 cmp_fn: MergeFuncLR(cmp_fn, PhantomData),
91 }
92}
93
94/// An iterator adaptor that merge-joins items from the two base iterators in ascending order.
95///
96/// See [`.merge_join_by()`](crate::Itertools::merge_join_by) for more information.
97pub type MergeJoinBy<I, J, F> =
98 MergeBy<I, J, MergeFuncLR<F, <F as FuncLR<<I as Iterator>::Item, <J as Iterator>::Item>>::T>>;
99
100#[derive(Clone, Debug)]
101pub struct MergeFuncLR<F, T>(F, PhantomData<T>);
102
103pub trait FuncLR<L, R> {
104 type T;
105}
106
107impl<L, R, T, F: FnMut(&L, &R) -> T> FuncLR<L, R> for F {
108 type T = T;
109}
110
111pub trait OrderingOrBool<L, R> {
112 type MergeResult;
113 fn left(left: L) -> Self::MergeResult;
114 fn right(right: R) -> Self::MergeResult;
115 // "merge" never returns (Some(...), Some(...), ...) so Option<Either<I::Item, J::Item>>
116 // is appealing but it is always followed by two put_backs, so we think the compiler is
117 // smart enough to optimize it. Or we could move put_backs into "merge".
118 fn merge(&mut self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult);
119 fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint;
120}
121
122impl<L, R, F: FnMut(&L, &R) -> Ordering> OrderingOrBool<L, R> for MergeFuncLR<F, Ordering> {
123 type MergeResult = EitherOrBoth<L, R>;
124 fn left(left: L) -> Self::MergeResult {
125 EitherOrBoth::Left(left)
126 }
127 fn right(right: R) -> Self::MergeResult {
128 EitherOrBoth::Right(right)
129 }
130 fn merge(&mut self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult) {
131 match self.0(&left, &right) {
132 Ordering::Equal => (None, None, EitherOrBoth::Both(left, right)),
133 Ordering::Less => (None, Some(right), EitherOrBoth::Left(left)),
134 Ordering::Greater => (Some(left), None, EitherOrBoth::Right(right)),
135 }
136 }
137 fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint {
138 let (a_lower, a_upper) = left;
139 let (b_lower, b_upper) = right;
140 let lower = ::std::cmp::max(a_lower, b_lower);
141 let upper = match (a_upper, b_upper) {
142 (Some(x), Some(y)) => x.checked_add(y),
143 _ => None,
144 };
145 (lower, upper)
146 }
147}
148
149impl<L, R, F: FnMut(&L, &R) -> bool> OrderingOrBool<L, R> for MergeFuncLR<F, bool> {
150 type MergeResult = Either<L, R>;
151 fn left(left: L) -> Self::MergeResult {
152 Either::Left(left)
153 }
154 fn right(right: R) -> Self::MergeResult {
155 Either::Right(right)
156 }
157 fn merge(&mut self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult) {
158 if self.0(&left, &right) {
159 (None, Some(right), Either::Left(left))
160 } else {
161 (Some(left), None, Either::Right(right))
162 }
163 }
164 fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint {
165 // Not ExactSizeIterator because size may be larger than usize
166 size_hint::add(a:left, b:right)
167 }
168}
169
170impl<T, F: FnMut(&T, &T) -> bool> OrderingOrBool<T, T> for F {
171 type MergeResult = T;
172 fn left(left: T) -> Self::MergeResult {
173 left
174 }
175 fn right(right: T) -> Self::MergeResult {
176 right
177 }
178 fn merge(&mut self, left: T, right: T) -> (Option<T>, Option<T>, Self::MergeResult) {
179 if self(&left, &right) {
180 (None, Some(right), left)
181 } else {
182 (Some(left), None, right)
183 }
184 }
185 fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint {
186 // Not ExactSizeIterator because size may be larger than usize
187 size_hint::add(a:left, b:right)
188 }
189}
190
191impl<T: PartialOrd> OrderingOrBool<T, T> for MergeLte {
192 type MergeResult = T;
193 fn left(left: T) -> Self::MergeResult {
194 left
195 }
196 fn right(right: T) -> Self::MergeResult {
197 right
198 }
199 fn merge(&mut self, left: T, right: T) -> (Option<T>, Option<T>, Self::MergeResult) {
200 if left <= right {
201 (None, Some(right), left)
202 } else {
203 (Some(left), None, right)
204 }
205 }
206 fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint {
207 // Not ExactSizeIterator because size may be larger than usize
208 size_hint::add(a:left, b:right)
209 }
210}
211
212impl<I, J, F> Clone for MergeBy<I, J, F>
213where
214 I: Iterator,
215 J: Iterator,
216 PutBack<Fuse<I>>: Clone,
217 PutBack<Fuse<J>>: Clone,
218 F: Clone,
219{
220 clone_fields!(left, right, cmp_fn);
221}
222
223impl<I, J, F> fmt::Debug for MergeBy<I, J, F>
224where
225 I: Iterator + fmt::Debug,
226 I::Item: fmt::Debug,
227 J: Iterator + fmt::Debug,
228 J::Item: fmt::Debug,
229{
230 debug_fmt_fields!(MergeBy, left, right);
231}
232
233impl<I, J, F> Iterator for MergeBy<I, J, F>
234where
235 I: Iterator,
236 J: Iterator,
237 F: OrderingOrBool<I::Item, J::Item>,
238{
239 type Item = F::MergeResult;
240
241 fn next(&mut self) -> Option<Self::Item> {
242 match (self.left.next(), self.right.next()) {
243 (None, None) => None,
244 (Some(left), None) => Some(F::left(left)),
245 (None, Some(right)) => Some(F::right(right)),
246 (Some(left), Some(right)) => {
247 let (left, right, next) = self.cmp_fn.merge(left, right);
248 if let Some(left) = left {
249 self.left.put_back(left);
250 }
251 if let Some(right) = right {
252 self.right.put_back(right);
253 }
254 Some(next)
255 }
256 }
257 }
258
259 fn size_hint(&self) -> SizeHint {
260 F::size_hint(self.left.size_hint(), self.right.size_hint())
261 }
262
263 fn count(mut self) -> usize {
264 let mut count = 0;
265 loop {
266 match (self.left.next(), self.right.next()) {
267 (None, None) => break count,
268 (Some(_left), None) => break count + 1 + self.left.into_parts().1.count(),
269 (None, Some(_right)) => break count + 1 + self.right.into_parts().1.count(),
270 (Some(left), Some(right)) => {
271 count += 1;
272 let (left, right, _) = self.cmp_fn.merge(left, right);
273 if let Some(left) = left {
274 self.left.put_back(left);
275 }
276 if let Some(right) = right {
277 self.right.put_back(right);
278 }
279 }
280 }
281 }
282 }
283
284 fn last(mut self) -> Option<Self::Item> {
285 let mut previous_element = None;
286 loop {
287 match (self.left.next(), self.right.next()) {
288 (None, None) => break previous_element,
289 (Some(left), None) => {
290 break Some(F::left(self.left.into_parts().1.last().unwrap_or(left)))
291 }
292 (None, Some(right)) => {
293 break Some(F::right(self.right.into_parts().1.last().unwrap_or(right)))
294 }
295 (Some(left), Some(right)) => {
296 let (left, right, elem) = self.cmp_fn.merge(left, right);
297 if let Some(left) = left {
298 self.left.put_back(left);
299 }
300 if let Some(right) = right {
301 self.right.put_back(right);
302 }
303 previous_element = Some(elem);
304 }
305 }
306 }
307 }
308
309 fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
310 loop {
311 if n == 0 {
312 break self.next();
313 }
314 n -= 1;
315 match (self.left.next(), self.right.next()) {
316 (None, None) => break None,
317 (Some(_left), None) => break self.left.nth(n).map(F::left),
318 (None, Some(_right)) => break self.right.nth(n).map(F::right),
319 (Some(left), Some(right)) => {
320 let (left, right, _) = self.cmp_fn.merge(left, right);
321 if let Some(left) = left {
322 self.left.put_back(left);
323 }
324 if let Some(right) = right {
325 self.right.put_back(right);
326 }
327 }
328 }
329 }
330 }
331}
332
333impl<I, J, F> FusedIterator for MergeBy<I, J, F>
334where
335 I: Iterator,
336 J: Iterator,
337 F: OrderingOrBool<I::Item, J::Item>,
338{
339}
340