Skip to content

Commit

Permalink
Add a new type for the factored iterator
Browse files Browse the repository at this point in the history
  • Loading branch information
cuviper committed Feb 10, 2024
1 parent 356d9b6 commit c1ec796
Show file tree
Hide file tree
Showing 2 changed files with 169 additions and 10 deletions.
178 changes: 168 additions & 10 deletions src/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,23 +84,25 @@ impl<L, R> Either<L, R> {
/// ```
// TODO(MSRV): doc(alias) was stabilized in Rust 1.48
// #[doc(alias = "transpose")]
pub fn factor_iter(
self,
) -> Either<
core::iter::Map<L::IntoIter, impl Fn(L::Item) -> Either<L::Item, R::Item>>,
core::iter::Map<R::IntoIter, impl Fn(R::Item) -> Either<L::Item, R::Item>>,
>
pub fn factor_iter(self) -> IterEither<L::IntoIter, R::IntoIter>
where
L: IntoIterator,
R: IntoIterator,
{
self.map_either(
|l| l.into_iter().map(Either::Left),
|r| r.into_iter().map(Either::Right),
)
IterEither {
inner: self.map_either(L::into_iter, R::into_iter),
}
}
}

/// Iterator that maps left or right iterators to corresponding `Either`-wrapped items.
///
/// This struct is created by the [`Either::factor_iter`] method.
#[derive(Clone, Debug)]
pub struct IterEither<L, R> {
inner: Either<L, R>,
}

impl<L, R, A> Extend<A> for Either<L, R>
where
L: Extend<A>,
Expand Down Expand Up @@ -252,3 +254,159 @@ where
R: iter::FusedIterator<Item = L::Item>,
{
}

macro_rules! map_either {
($value:expr, $pattern:pat => $result:expr) => {
match $value {
Left($pattern) => Left($result),
Right($pattern) => Right($result),
}
};
}

macro_rules! wrap_either {
($value:expr => $( $tail:tt )*) => {
match $value {
Left(inner) => inner.map(Left) $($tail)*,
Right(inner) => inner.map(Right) $($tail)*,
}
};
}

impl<L, R> Iterator for IterEither<L, R>
where
L: Iterator,
R: Iterator,
{
type Item = Either<L::Item, R::Item>;

fn next(&mut self) -> Option<Self::Item> {
Some(map_either!(self.inner, ref mut inner => inner.next()?))
}

fn size_hint(&self) -> (usize, Option<usize>) {
for_both!(self.inner, ref inner => inner.size_hint())
}

fn fold<Acc, G>(self, init: Acc, f: G) -> Acc
where
G: FnMut(Acc, Self::Item) -> Acc,
{
wrap_either!(self.inner => .fold(init, f))
}

fn for_each<F>(self, f: F)
where
F: FnMut(Self::Item),
{
wrap_either!(self.inner => .for_each(f))
}

fn count(self) -> usize {
for_both!(self.inner, inner => inner.count())
}

fn last(self) -> Option<Self::Item> {
Some(map_either!(self.inner, inner => inner.last()?))
}

fn nth(&mut self, n: usize) -> Option<Self::Item> {
Some(map_either!(self.inner, ref mut inner => inner.nth(n)?))
}

fn collect<B>(self) -> B
where
B: iter::FromIterator<Self::Item>,
{
wrap_either!(self.inner => .collect())
}

fn partition<B, F>(self, f: F) -> (B, B)
where
B: Default + Extend<Self::Item>,
F: FnMut(&Self::Item) -> bool,
{
wrap_either!(self.inner => .partition(f))
}

fn all<F>(&mut self, f: F) -> bool
where
F: FnMut(Self::Item) -> bool,
{
wrap_either!(&mut self.inner => .all(f))
}

fn any<F>(&mut self, f: F) -> bool
where
F: FnMut(Self::Item) -> bool,
{
wrap_either!(&mut self.inner => .any(f))
}

fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
where
P: FnMut(&Self::Item) -> bool,
{
wrap_either!(&mut self.inner => .find(predicate))
}

fn find_map<B, F>(&mut self, f: F) -> Option<B>
where
F: FnMut(Self::Item) -> Option<B>,
{
wrap_either!(&mut self.inner => .find_map(f))
}

fn position<P>(&mut self, predicate: P) -> Option<usize>
where
P: FnMut(Self::Item) -> bool,
{
wrap_either!(&mut self.inner => .position(predicate))
}
}

impl<L, R> DoubleEndedIterator for IterEither<L, R>
where
L: DoubleEndedIterator,
R: DoubleEndedIterator,
{
fn next_back(&mut self) -> Option<Self::Item> {
Some(map_either!(self.inner, ref mut inner => inner.next_back()?))
}

// TODO(MSRV): This was stabilized in Rust 1.37
// fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
// Some(map_either!(self.inner, ref mut inner => inner.nth_back(n)?))
// }

fn rfold<Acc, G>(self, init: Acc, f: G) -> Acc
where
G: FnMut(Acc, Self::Item) -> Acc,
{
wrap_either!(self.inner => .rfold(init, f))
}

fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
where
P: FnMut(&Self::Item) -> bool,
{
wrap_either!(&mut self.inner => .rfind(predicate))
}
}

impl<L, R> ExactSizeIterator for IterEither<L, R>
where
L: ExactSizeIterator,
R: ExactSizeIterator,
{
fn len(&self) -> usize {
for_both!(self.inner, ref inner => inner.len())
}
}

impl<L, R> iter::FusedIterator for IterEither<L, R>
where
L: iter::FusedIterator,
R: iter::FusedIterator,
{
}
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -929,6 +929,7 @@ impl<T> Either<T, T> {
}

mod iterator;
pub use iterator::IterEither;

/// Convert from `Result` to `Either` with `Ok => Right` and `Err => Left`.
impl<L, R> From<Result<R, L>> for Either<L, R> {
Expand Down

0 comments on commit c1ec796

Please sign in to comment.