1use std::cmp::Ordering;
2use std::iter::Fuse;
3use std::fmt;
4
5use either::Either;
6
7use super::adaptors::{PutBack, put_back};
8use crate::either_or_both::EitherOrBoth;
9use crate::size_hint::{self, SizeHint};
10#[cfg(doc)]
11use 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`].
16pub 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(iterable:left.into_iter().fuse()),
25 right: put_back(iterable: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"]
34pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
35 left: PutBack<Fuse<I>>,
36 right: PutBack<Fuse<J>>,
37 cmp_fn: F,
38}
39
40pub trait OrderingOrBool<L, R> {
41 type MergeResult;
42 fn left(left: L) -> Self::MergeResult;
43 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".
47 fn merge(self, left: L, right: R) -> (Option<L>, Option<R>, Self::MergeResult);
48 fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint;
49}
50
51impl<L, R> OrderingOrBool<L, R> for Ordering {
52 type MergeResult = EitherOrBoth<L, R>;
53 fn left(left: L) -> Self::MergeResult {
54 EitherOrBoth::Left(left)
55 }
56 fn right(right: R) -> Self::MergeResult {
57 EitherOrBoth::Right(right)
58 }
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 }
66 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
78impl<L, R> OrderingOrBool<L, R> for bool {
79 type MergeResult = Either<L, R>;
80 fn left(left: L) -> Self::MergeResult {
81 Either::Left(left)
82 }
83 fn right(right: R) -> Self::MergeResult {
84 Either::Right(right)
85 }
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 }
93 fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint {
94 // Not ExactSizeIterator because size may be larger than usize
95 size_hint::add(a:left, b:right)
96 }
97}
98
99impl<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
109impl<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
118impl<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
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
144 fn size_hint(&self) -> SizeHint {
145 T::size_hint(self.left.size_hint(), self.right.size_hint())
146 }
147
148 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
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
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