itertools/
merge_join.rs

1use std::cmp::Ordering;
2use std::iter::Fuse;
3use std::fmt;
4
5use super::adaptors::{PutBack, put_back};
6use crate::either_or_both::EitherOrBoth;
7#[cfg(doc)]
8use 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`].
13pub 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"]
30pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
31    left: PutBack<Fuse<I>>,
32    right: PutBack<Fuse<J>>,
33    cmp_fn: F
34}
35
36impl<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
46impl<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
55impl<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}
itertools/merge_join.rs

itertools/
merge_join.rs
