diff --git a/src/lib.rs b/src/lib.rs index 23b1bda01..340a04120 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1014,7 +1014,7 @@ pub trait Itertools : Iterator { /// /// The function can either return an `Ordering` variant or a boolean. /// - /// In the first case, + /// If `cmp_fn` returns `Ordering`, /// it chooses whether to pair elements based on the `Ordering` returned by the /// specified compare function. At any point, inspecting the tip of the /// iterators `I` and `J` as items `i` of type `I::Item` and `j` of type @@ -1031,16 +1031,16 @@ pub trait Itertools : Iterator { /// use itertools::Itertools; /// use itertools::EitherOrBoth::{Left, Right, Both}; /// - /// let multiples_of_2 = (0..10).step_by(2); - /// let multiples_of_3 = (0..10).step_by(3); + /// let a = vec![0, 2, 4, 6, 1].into_iter(); + /// let b = (0..10).step_by(3); /// /// itertools::assert_equal( - /// multiples_of_2.merge_join_by(multiples_of_3, |i, j| i.cmp(j)), - /// vec![Both(0, 0), Left(2), Right(3), Left(4), Both(6, 6), Left(8), Right(9)] + /// a.merge_join_by(b, |i, j| i.cmp(j)), + /// vec![Both(0, 0), Left(2), Right(3), Left(4), Both(6, 6), Left(1), Right(9)] /// ); /// ``` /// - /// In the second case, + /// If `cmp_fn` returns `bool`, /// it chooses whether to pair elements based on the boolean returned by the /// specified function. At any point, inspecting the tip of the /// iterators `I` and `J` as items `i` of type `I::Item` and `j` of type @@ -1051,22 +1051,26 @@ pub trait Itertools : Iterator { /// - Emit `Either::Right(j)` when `false`, /// and remove `j` from its source iterator /// + /// It is similar to the `Ordering` case if the first argument is considered + /// "less" than the second argument. + /// /// ``` /// use itertools::Itertools; /// use itertools::Either::{Left, Right}; /// - /// let multiples_of_2 = (0..10).step_by(2); - /// let multiples_of_3 = (0..10).step_by(3); + /// let a = vec![0, 2, 4, 6, 1].into_iter(); + /// let b = (0..10).step_by(3); /// /// itertools::assert_equal( - /// multiples_of_2.merge_join_by(multiples_of_3, |i, j| i <= j), - /// vec![Left(0), Right(0), Left(2), Right(3), Left(4), Left(6), Right(6), Left(8), Right(9)] + /// a.merge_join_by(b, |i, j| i <= j), + /// vec![Left(0), Right(0), Left(2), Right(3), Left(4), Left(6), Left(1), Right(6), Right(9)] /// ); /// ``` #[inline] fn merge_join_by(self, other: J, cmp_fn: F) -> MergeJoinBy where J: IntoIterator, F: FnMut(&Self::Item, &J::Item) -> T, + T: merge_join::OrderingOrBool, Self: Sized { merge_join_by(self, other, cmp_fn) diff --git a/src/merge_join.rs b/src/merge_join.rs index 15b17ce79..4c7e35cdc 100644 --- a/src/merge_join.rs +++ b/src/merge_join.rs @@ -18,6 +18,7 @@ pub fn merge_join_by(left: I, right: J, cmp_fn: F) where I: IntoIterator, J: IntoIterator, F: FnMut(&I::Item, &J::Item) -> T, + T: OrderingOrBool, { MergeJoinBy { left: put_back(left.into_iter().fuse()), @@ -40,6 +41,9 @@ pub trait OrderingOrBool { type MergeResult; fn left(left: I) -> Self::MergeResult; fn right(right: J) -> Self::MergeResult; + // "merge" never returns (Some(...), Some(...), ...) so Option> + // is appealing but it is always followed by two put_backs, so we think the compiler is + // smart enough to optimize it. Or we could move put_backs into "merge". fn merge(self, left: I, right: J) -> (Option, Option, Self::MergeResult); fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint; } @@ -179,13 +183,13 @@ impl Iterator for MergeJoinBy } (Some(left), Some(right)) => { let (left, right, elem) = (self.cmp_fn)(&left, &right).merge(left, right); - previous_element = Some(elem); if let Some(left) = left { self.left.put_back(left); } if let Some(right) = right { self.right.put_back(right); } + previous_element = Some(elem); } } } diff --git a/tests/quick.rs b/tests/quick.rs index 0adcf1ad7..fbb731f26 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -829,6 +829,33 @@ quickcheck! { } } +quickcheck! { + fn merge_join_by_ordering_vs_bool(a: Vec, b: Vec) -> bool { + use either::Either; + use itertools::EitherOrBoth; + use itertools::free::merge_join_by; + let mut has_equal = false; + let it_ord = merge_join_by(a.clone(), b.clone(), Ord::cmp).flat_map(|v| match v { + EitherOrBoth::Both(l, r) => { + has_equal = true; + vec![l, r] + } + EitherOrBoth::Left(l) => vec![l], + EitherOrBoth::Right(r) => vec![r], + }); + let it_bool = merge_join_by(a, b, PartialOrd::le).map(Either::into_inner); + itertools::equal(it_ord, it_bool) || has_equal + } + fn merge_join_by_bool_unwrapped_is_merge_by(a: Vec, b: Vec) -> bool { + use either::Either; + use itertools::Itertools; + use itertools::free::merge_join_by; + let it = a.clone().into_iter().merge_by(b.clone(), PartialOrd::ge); + let it_join = merge_join_by(a, b, PartialOrd::ge).map(Either::into_inner); + itertools::equal(it, it_join) + } +} + quickcheck! { fn size_tee(a: Vec) -> bool { let (mut t1, mut t2) = a.iter().tee();