itertools/
merge_join.rs

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(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"]
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(left, 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}
itertools/merge_join.rs

itertools/
merge_join.rs
