From 2f45b173d68215aa4272aaeac0ac6f09d3d3c72e Mon Sep 17 00:00:00 2001
From: Philippe-Cholet <phcholet@orange.fr>
Date: Sun, 11 Jun 2023 09:43:21 +0200
Subject: [PATCH] `MergeBy` wraps `MergeJoinBy`

---
 src/adaptors/merge.rs | 95 +++++++++++++------------------------------
 1 file changed, 29 insertions(+), 66 deletions(-)

diff --git a/src/adaptors/merge.rs b/src/adaptors/merge.rs
index 4d42f30b2..fb1e99767 100644
--- a/src/adaptors/merge.rs
+++ b/src/adaptors/merge.rs
@@ -1,20 +1,8 @@
 use std::fmt;
-use std::iter::{Peekable, FusedIterator};
 
-use crate::size_hint;
+use either::Either;
 
-pub trait MergePredicate<T> {
-    fn merge_pred(&mut self, a: &T, b: &T) -> bool;
-}
-
-#[derive(Clone, Debug)]
-pub struct MergeLte;
-
-impl<T: PartialOrd> MergePredicate<T> for MergeLte {
-    fn merge_pred(&mut self, a: &T, b: &T) -> bool {
-        a <= b
-    }
-}
+use crate::merge_join::{MergeJoinBy, merge_join_by};
 
 /// An iterator adaptor that merges the two base iterators in ascending order.
 /// If both base iterators are sorted (ascending), the result is sorted.
@@ -22,7 +10,7 @@ impl<T: PartialOrd> MergePredicate<T> for MergeLte {
 /// Iterator element type is `I::Item`.
 ///
 /// See [`.merge()`](crate::Itertools::merge_by) for more information.
-pub type Merge<I, J> = MergeBy<I, J, MergeLte>;
+pub type Merge<I, J> = MergeBy<I, J, fn(&<I as Iterator>::Item, &<J as Iterator>::Item) -> bool>;
 
 /// Create an iterator that merges elements in `i` and `j`.
 ///
@@ -40,7 +28,7 @@ pub fn merge<I, J>(i: I, j: J) -> Merge<<I as IntoIterator>::IntoIter, <J as Int
           J: IntoIterator<Item = I::Item>,
           I::Item: PartialOrd
 {
-    merge_by_new(i, j, MergeLte)
+    merge_by_new(i, j, I::Item::le)
 }
 
 /// An iterator adaptor that merges the two base iterators in ascending order.
@@ -52,89 +40,64 @@ pub fn merge<I, J>(i: I, j: J) -> Merge<<I as IntoIterator>::IntoIter, <J as Int
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 pub struct MergeBy<I, J, F>
     where I: Iterator,
-          J: Iterator<Item = I::Item>
+          J: Iterator<Item = I::Item>,
+          F: FnMut(&I::Item, &I::Item) -> bool,
 {
-    a: Peekable<I>,
-    b: Peekable<J>,
-    fused: Option<bool>,
-    cmp: F,
+    inner: MergeJoinBy<I, J, F, bool>,
 }
 
 impl<I, J, F> fmt::Debug for MergeBy<I, J, F>
     where I: Iterator + fmt::Debug, J: Iterator<Item = I::Item> + fmt::Debug,
           I::Item: fmt::Debug,
+          F: FnMut(&I::Item, &I::Item) -> bool,
 {
-    debug_fmt_fields!(MergeBy, a, b);
-}
-
-impl<T, F: FnMut(&T, &T) -> bool> MergePredicate<T> for F {
-    fn merge_pred(&mut self, a: &T, b: &T) -> bool {
-        self(a, b)
-    }
+    debug_fmt_fields!(MergeBy, inner);
 }
 
 /// Create a `MergeBy` iterator.
 pub fn merge_by_new<I, J, F>(a: I, b: J, cmp: F) -> MergeBy<I::IntoIter, J::IntoIter, F>
     where I: IntoIterator,
           J: IntoIterator<Item = I::Item>,
-          F: MergePredicate<I::Item>,
+          F: FnMut(&I::Item, &I::Item) -> bool,
 {
     MergeBy {
-        a: a.into_iter().peekable(),
-        b: b.into_iter().peekable(),
-        fused: None,
-        cmp,
+        inner: merge_join_by(a.into_iter(), b.into_iter(), cmp),
     }
 }
 
 impl<I, J, F> Clone for MergeBy<I, J, F>
     where I: Iterator,
           J: Iterator<Item = I::Item>,
-          Peekable<I>: Clone,
-          Peekable<J>: Clone,
-          F: Clone
+          F: FnMut(&I::Item, &I::Item) -> bool,
+          MergeJoinBy<I, J, F, bool>: Clone,
 {
-    clone_fields!(a, b, fused, cmp);
+    clone_fields!(inner);
 }
 
 impl<I, J, F> Iterator for MergeBy<I, J, F>
     where I: Iterator,
           J: Iterator<Item = I::Item>,
-          F: MergePredicate<I::Item>
+          F: FnMut(&I::Item, &I::Item) -> bool,
 {
     type Item = I::Item;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let less_than = match self.fused {
-            Some(lt) => lt,
-            None => match (self.a.peek(), self.b.peek()) {
-                (Some(a), Some(b)) => self.cmp.merge_pred(a, b),
-                (Some(_), None) => {
-                    self.fused = Some(true);
-                    true
-                }
-                (None, Some(_)) => {
-                    self.fused = Some(false);
-                    false
-                }
-                (None, None) => return None,
-            }
-        };
-        if less_than {
-            self.a.next()
-        } else {
-            self.b.next()
-        }
+        self.inner.next().map(Either::into_inner)
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        // Not ExactSizeIterator because size may be larger than usize
-        size_hint::add(self.a.size_hint(), self.b.size_hint())
+        self.inner.size_hint()
+    }
+
+    fn count(self) -> usize {
+        self.inner.count()
     }
-}
 
-impl<I, J, F> FusedIterator for MergeBy<I, J, F>
-    where I: FusedIterator,
-          J: FusedIterator<Item = I::Item>,
-          F: MergePredicate<I::Item>
-{}
+    fn last(self) -> Option<Self::Item> {
+        self.inner.last().map(Either::into_inner)
+    }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        self.inner.nth(n).map(Either::into_inner)
+    }
+}