Skip to content

Commit

Permalink
MergeBy wraps MergeJoinBy
Browse files Browse the repository at this point in the history
  • Loading branch information
Philippe-Cholet committed Jun 11, 2023
1 parent d853808 commit 2f45b17
Showing 1 changed file with 29 additions and 66 deletions.
95 changes: 29 additions & 66 deletions src/adaptors/merge.rs
Original file line number Diff line number Diff line change
@@ -1,28 +1,16 @@
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.
///
/// 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`.
///
Expand All @@ -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.
Expand All @@ -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)
}
}

0 comments on commit 2f45b17

Please sign in to comment.