1 | use std::cmp::Ordering; |
2 | use std::iter::Fuse; |
3 | use std::fmt; |
4 | |
5 | use super::adaptors::{PutBack, put_back}; |
6 | use crate::either_or_both::EitherOrBoth; |
7 | #[cfg (doc)] |
8 | use crate::Itertools; |
9 | |
10 | /// Return an iterator adaptor that merge-joins items from the two base iterators in ascending order. |
11 | /// |
12 | /// [`IntoIterator`] enabled version of [`Itertools::merge_join_by`]. |
13 | pub fn merge_join_by<I, J, F>(left: I, right: J, cmp_fn: F) |
14 | -> MergeJoinBy<I::IntoIter, J::IntoIter, F> |
15 | where I: IntoIterator, |
16 | J: IntoIterator, |
17 | F: FnMut(&I::Item, &J::Item) -> Ordering |
18 | { |
19 | MergeJoinBy { |
20 | left: put_back(left.into_iter().fuse()), |
21 | right: put_back(right.into_iter().fuse()), |
22 | cmp_fn, |
23 | } |
24 | } |
25 | |
26 | /// An iterator adaptor that merge-joins items from the two base iterators in ascending order. |
27 | /// |
28 | /// See [`.merge_join_by()`](crate::Itertools::merge_join_by) for more information. |
29 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed" ] |
30 | pub struct MergeJoinBy<I: Iterator, J: Iterator, F> { |
31 | left: PutBack<Fuse<I>>, |
32 | right: PutBack<Fuse<J>>, |
33 | cmp_fn: F |
34 | } |
35 | |
36 | impl<I, J, F> Clone for MergeJoinBy<I, J, F> |
37 | where I: Iterator, |
38 | J: Iterator, |
39 | PutBack<Fuse<I>>: Clone, |
40 | PutBack<Fuse<J>>: Clone, |
41 | F: Clone, |
42 | { |
43 | clone_fields!(left, right, cmp_fn); |
44 | } |
45 | |
46 | impl<I, J, F> fmt::Debug for MergeJoinBy<I, J, F> |
47 | where I: Iterator + fmt::Debug, |
48 | I::Item: fmt::Debug, |
49 | J: Iterator + fmt::Debug, |
50 | J::Item: fmt::Debug, |
51 | { |
52 | debug_fmt_fields!(MergeJoinBy, left, right); |
53 | } |
54 | |
55 | impl<I, J, F> Iterator for MergeJoinBy<I, J, F> |
56 | where I: Iterator, |
57 | J: Iterator, |
58 | F: FnMut(&I::Item, &J::Item) -> Ordering |
59 | { |
60 | type Item = EitherOrBoth<I::Item, J::Item>; |
61 | |
62 | fn next(&mut self) -> Option<Self::Item> { |
63 | match (self.left.next(), self.right.next()) { |
64 | (None, None) => None, |
65 | (Some(left), None) => |
66 | Some(EitherOrBoth::Left(left)), |
67 | (None, Some(right)) => |
68 | Some(EitherOrBoth::Right(right)), |
69 | (Some(left), Some(right)) => { |
70 | match (self.cmp_fn)(&left, &right) { |
71 | Ordering::Equal => |
72 | Some(EitherOrBoth::Both(left, right)), |
73 | Ordering::Less => { |
74 | self.right.put_back(right); |
75 | Some(EitherOrBoth::Left(left)) |
76 | }, |
77 | Ordering::Greater => { |
78 | self.left.put_back(left); |
79 | Some(EitherOrBoth::Right(right)) |
80 | } |
81 | } |
82 | } |
83 | } |
84 | } |
85 | |
86 | fn size_hint(&self) -> (usize, Option<usize>) { |
87 | let (a_lower, a_upper) = self.left.size_hint(); |
88 | let (b_lower, b_upper) = self.right.size_hint(); |
89 | |
90 | let lower = ::std::cmp::max(a_lower, b_lower); |
91 | |
92 | let upper = match (a_upper, b_upper) { |
93 | (Some(x), Some(y)) => x.checked_add(y), |
94 | _ => None, |
95 | }; |
96 | |
97 | (lower, upper) |
98 | } |
99 | |
100 | fn count(mut self) -> usize { |
101 | let mut count = 0; |
102 | loop { |
103 | match (self.left.next(), self.right.next()) { |
104 | (None, None) => break count, |
105 | (Some(_left), None) => break count + 1 + self.left.into_parts().1.count(), |
106 | (None, Some(_right)) => break count + 1 + self.right.into_parts().1.count(), |
107 | (Some(left), Some(right)) => { |
108 | count += 1; |
109 | match (self.cmp_fn)(&left, &right) { |
110 | Ordering::Equal => {} |
111 | Ordering::Less => self.right.put_back(right), |
112 | Ordering::Greater => self.left.put_back(left), |
113 | } |
114 | } |
115 | } |
116 | } |
117 | } |
118 | |
119 | fn last(mut self) -> Option<Self::Item> { |
120 | let mut previous_element = None; |
121 | loop { |
122 | match (self.left.next(), self.right.next()) { |
123 | (None, None) => break previous_element, |
124 | (Some(left), None) => { |
125 | break Some(EitherOrBoth::Left( |
126 | self.left.into_parts().1.last().unwrap_or(left), |
127 | )) |
128 | } |
129 | (None, Some(right)) => { |
130 | break Some(EitherOrBoth::Right( |
131 | self.right.into_parts().1.last().unwrap_or(right), |
132 | )) |
133 | } |
134 | (Some(left), Some(right)) => { |
135 | previous_element = match (self.cmp_fn)(&left, &right) { |
136 | Ordering::Equal => Some(EitherOrBoth::Both(left, right)), |
137 | Ordering::Less => { |
138 | self.right.put_back(right); |
139 | Some(EitherOrBoth::Left(left)) |
140 | } |
141 | Ordering::Greater => { |
142 | self.left.put_back(left); |
143 | Some(EitherOrBoth::Right(right)) |
144 | } |
145 | } |
146 | } |
147 | } |
148 | } |
149 | } |
150 | |
151 | fn nth(&mut self, mut n: usize) -> Option<Self::Item> { |
152 | loop { |
153 | if n == 0 { |
154 | break self.next(); |
155 | } |
156 | n -= 1; |
157 | match (self.left.next(), self.right.next()) { |
158 | (None, None) => break None, |
159 | (Some(_left), None) => break self.left.nth(n).map(EitherOrBoth::Left), |
160 | (None, Some(_right)) => break self.right.nth(n).map(EitherOrBoth::Right), |
161 | (Some(left), Some(right)) => match (self.cmp_fn)(&left, &right) { |
162 | Ordering::Equal => {} |
163 | Ordering::Less => self.right.put_back(right), |
164 | Ordering::Greater => self.left.put_back(left), |
165 | }, |
166 | } |
167 | } |
168 | } |
169 | } |
170 | |