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