merge_join_by: Ordering vs bool

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<J, F, T>(self, other: J, cmp_fn: F) -> MergeJoinBy<Self, J::IntoIter, F>
         where J: IntoIterator,
               F: FnMut(&Self::Item, &J::Item) -> T,
+              T: merge_join::OrderingOrBool<Self::Item, J::Item>,
               Self: Sized
     {
         merge_join_by(self, other, cmp_fn)

src/merge_join.rs

@@ -18,6 +18,7 @@ pub fn merge_join_by<I, J, F, T>(left: I, right: J, cmp_fn: F)
     where I: IntoIterator,
           J: IntoIterator,
           F: FnMut(&I::Item, &J::Item) -> T,
+          T: OrderingOrBool<I::Item, J::Item>,
 {
     MergeJoinBy {
         left: put_back(left.into_iter().fuse()),
@@ -40,6 +41,9 @@ pub trait OrderingOrBool<I, J> {
     type MergeResult;
     fn left(left: I) -> Self::MergeResult;
     fn right(right: J) -> Self::MergeResult;
+    // "merge" never returns (Some(...), Some(...), ...) so Option<Either<I::Item, J::Item>>
+    // 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<I>, Option<J>, Self::MergeResult);
     fn size_hint(left: SizeHint, right: SizeHint) -> SizeHint;
 }
@@ -179,13 +183,13 @@ impl<I, J, F, T> Iterator for MergeJoinBy<I, J, F>
                 }
                 (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);
                 }
             }
         }

tests/quick.rs

@@ -829,6 +829,33 @@ quickcheck! {
     }
 }
 
+quickcheck! {
+    fn merge_join_by_ordering_vs_bool(a: Vec<u8>, b: Vec<u8>) -> 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<u8>, b: Vec<u8>) -> 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<u8>) -> bool {
         let (mut t1, mut t2) = a.iter().tee();