1 use std::cmp::Ordering;
2 use std::iter::Fuse;
3 use std::fmt;
4 
5 use either::Either;
6 
7 use super::adaptors::{PutBack, put_back};
8 use crate::either_or_both::EitherOrBoth;
9 use crate::size_hint::{self, SizeHint};
10 #[cfg(doc)]
11 use crate::Itertools;
12 
13 /// Return an iterator adaptor that merge-joins items from the two base iterators in ascending order.
14 ///
15 /// [`IntoIterator`] enabled version of [`Itertools::merge_join_by`].
merge_join_by<I, J, F, T>(left: I, right: J, cmp_fn: F) -> MergeJoinBy<I::IntoIter, J::IntoIter, F> where I: IntoIterator, J: IntoIterator, F: FnMut(&I::Item, &J::Item) -> T, T: OrderingOrBool<I::Item, J::Item>,16 pub fn merge_join_by<I, J, F, T>(left: I, right: J, cmp_fn: F)
17     -> MergeJoinBy<I::IntoIter, J::IntoIter, F>
18     where I: IntoIterator,
19           J: IntoIterator,
20           F: FnMut(&I::Item, &J::Item) -> T,
21           T: OrderingOrBool<I::Item, J::Item>,
22 {
23     MergeJoinBy {
24         left: put_back(left.into_iter().fuse()),
25         right: put_back(right.into_iter().fuse()),
26         cmp_fn,
27     }
28 }
29 
30 /// An iterator adaptor that merge-joins items from the two base iterators in ascending order.
31 ///
32 /// See [`.merge_join_by()`](crate::Itertools::merge_join_by) for more information.
33 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
34 pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
35     left: PutBack<Fuse<I>>,
36     right: PutBack<Fuse<J>>,
37     cmp_fn: F,
38 }
39 
40 pub trait OrderingOrBool<L, R> {
41     type MergeResult;
left(left: L) -> Self::MergeResult42     fn left(left: L) -> Self::MergeResult;
right(right: R) -> Self::MergeResult43     fn right(right: R) -> Self::MergeResult;
44     // "merge" never returns (Some(...), Some(...), ...) so Option<Either<I::Item, J::Item>>
45     // is appealing but it is always followed by two put_backs, so we think the compiler is
46     // smart enough to optimize it. Or we could move put_backs into "merge".
merge(self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult)47     fn merge(self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult);
size_hint(left: SizeHint, right: SizeHint) -> SizeHint48     fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint;
49 }
50 
51 impl<L, R> OrderingOrBool<L, R> for Ordering {
52     type MergeResult = EitherOrBoth<L, R>;
left(left: L) -> Self::MergeResult53     fn left(left: L) -> Self::MergeResult {
54         EitherOrBoth::Left(left)
55     }
right(right: R) -> Self::MergeResult56     fn right(right: R) -> Self::MergeResult {
57         EitherOrBoth::Right(right)
58     }
merge(self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult)59     fn merge(self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult) {
60         match self {
61             Ordering::Equal => (None, None, EitherOrBoth::Both(left, right)),
62             Ordering::Less => (None, Some(right), EitherOrBoth::Left(left)),
63             Ordering::Greater => (Some(left), None, EitherOrBoth::Right(right)),
64         }
65     }
size_hint(left: SizeHint, right: SizeHint) -> SizeHint66     fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint {
67         let (a_lower, a_upper) = left;
68         let (b_lower, b_upper) = right;
69         let lower = ::std::cmp::max(a_lower, b_lower);
70         let upper = match (a_upper, b_upper) {
71             (Some(x), Some(y)) => x.checked_add(y),
72             _ => None,
73         };
74         (lower, upper)
75     }
76 }
77 
78 impl<L, R> OrderingOrBool<L, R> for bool {
79     type MergeResult = Either<L, R>;
left(left: L) -> Self::MergeResult80     fn left(left: L) -> Self::MergeResult {
81         Either::Left(left)
82     }
right(right: R) -> Self::MergeResult83     fn right(right: R) -> Self::MergeResult {
84         Either::Right(right)
85     }
merge(self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult)86     fn merge(self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult) {
87         if self {
88             (None, Some(right), Either::Left(left))
89         } else {
90             (Some(left), None, Either::Right(right))
91         }
92     }
size_hint(left: SizeHint, right: SizeHint) -> SizeHint93     fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint {
94         // Not ExactSizeIterator because size may be larger than usize
95         size_hint::add(left, right)
96     }
97 }
98 
99 impl<I, J, F> Clone for MergeJoinBy<I, J, F>
100     where I: Iterator,
101           J: Iterator,
102           PutBack<Fuse<I>>: Clone,
103           PutBack<Fuse<J>>: Clone,
104           F: Clone,
105 {
106     clone_fields!(left, right, cmp_fn);
107 }
108 
109 impl<I, J, F> fmt::Debug for MergeJoinBy<I, J, F>
110     where I: Iterator + fmt::Debug,
111           I::Item: fmt::Debug,
112           J: Iterator + fmt::Debug,
113           J::Item: fmt::Debug,
114 {
115     debug_fmt_fields!(MergeJoinBy, left, right);
116 }
117 
118 impl<I, J, F, T> Iterator for MergeJoinBy<I, J, F>
119     where I: Iterator,
120           J: Iterator,
121           F: FnMut(&I::Item, &J::Item) -> T,
122           T: OrderingOrBool<I::Item, J::Item>,
123 {
124     type Item = T::MergeResult;
125 
next(&mut self) -> Option<Self::Item>126     fn next(&mut self) -> Option<Self::Item> {
127         match (self.left.next(), self.right.next()) {
128             (None, None) => None,
129             (Some(left), None) => Some(T::left(left)),
130             (None, Some(right)) => Some(T::right(right)),
131             (Some(left), Some(right)) => {
132                 let (left, right, next) = (self.cmp_fn)(&left, &right).merge(left, right);
133                 if let Some(left) = left {
134                     self.left.put_back(left);
135                 }
136                 if let Some(right) = right {
137                     self.right.put_back(right);
138                 }
139                 Some(next)
140             }
141         }
142     }
143 
size_hint(&self) -> SizeHint144     fn size_hint(&self) -> SizeHint {
145         T::size_hint(self.left.size_hint(), self.right.size_hint())
146     }
147 
count(mut self) -> usize148     fn count(mut self) -> usize {
149         let mut count = 0;
150         loop {
151             match (self.left.next(), self.right.next()) {
152                 (None, None) => break count,
153                 (Some(_left), None) => break count + 1 + self.left.into_parts().1.count(),
154                 (None, Some(_right)) => break count + 1 + self.right.into_parts().1.count(),
155                 (Some(left), Some(right)) => {
156                     count += 1;
157                     let (left, right, _) = (self.cmp_fn)(&left, &right).merge(left, right);
158                     if let Some(left) = left {
159                         self.left.put_back(left);
160                     }
161                     if let Some(right) = right {
162                         self.right.put_back(right);
163                     }
164                 }
165             }
166         }
167     }
168 
last(mut self) -> Option<Self::Item>169     fn last(mut self) -> Option<Self::Item> {
170         let mut previous_element = None;
171         loop {
172             match (self.left.next(), self.right.next()) {
173                 (None, None) => break previous_element,
174                 (Some(left), None) => {
175                     break Some(T::left(
176                         self.left.into_parts().1.last().unwrap_or(left),
177                     ))
178                 }
179                 (None, Some(right)) => {
180                     break Some(T::right(
181                         self.right.into_parts().1.last().unwrap_or(right),
182                     ))
183                 }
184                 (Some(left), Some(right)) => {
185                     let (left, right, elem) = (self.cmp_fn)(&left, &right).merge(left, right);
186                     if let Some(left) = left {
187                         self.left.put_back(left);
188                     }
189                     if let Some(right) = right {
190                         self.right.put_back(right);
191                     }
192                     previous_element = Some(elem);
193                 }
194             }
195         }
196     }
197 
nth(&mut self, mut n: usize) -> Option<Self::Item>198     fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
199         loop {
200             if n == 0 {
201                 break self.next();
202             }
203             n -= 1;
204             match (self.left.next(), self.right.next()) {
205                 (None, None) => break None,
206                 (Some(_left), None) => break self.left.nth(n).map(T::left),
207                 (None, Some(_right)) => break self.right.nth(n).map(T::right),
208                 (Some(left), Some(right)) => {
209                     let (left, right, _) = (self.cmp_fn)(&left, &right).merge(left, right);
210                     if let Some(left) = left {
211                         self.left.put_back(left);
212                     }
213                     if let Some(right) = right {
214                         self.right.put_back(right);
215                     }
216                 }
217             }
218         }
219     }
220 }
221