diff --git a/src/lib.rs b/src/lib.rs index 36ddef6cc..9a5b731ff 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -123,6 +123,7 @@ pub mod structs { pub use crate::pad_tail::PadUsing; #[cfg(feature = "use_alloc")] pub use crate::peek_nth::PeekNth; + pub use crate::peeking_map_while::PeekingMapWhile; pub use crate::peeking_take_while::PeekingTakeWhile; #[cfg(feature = "use_alloc")] pub use crate::permutations::Permutations; @@ -163,6 +164,7 @@ pub use crate::diff::Diff; #[cfg(feature = "use_alloc")] pub use crate::kmerge_impl::kmerge_by; pub use crate::minmax::MinMaxResult; +pub use crate::peeking_map_while::PeekingMapNext; pub use crate::peeking_take_while::PeekingNext; pub use crate::process_results_impl::process_results; pub use crate::repeatn::repeat_n; @@ -216,6 +218,7 @@ mod next_array; mod pad_tail; #[cfg(feature = "use_alloc")] mod peek_nth; +mod peeking_map_while; mod peeking_take_while; #[cfg(feature = "use_alloc")] mod permutations; @@ -1540,8 +1543,8 @@ pub trait Itertools: Iterator { /// Return an iterator adaptor that borrows from this iterator and /// takes items while the closure `accept` returns `true`. /// - /// This adaptor can only be used on iterators that implement `PeekingNext` - /// like `.peekable()`, `put_back` and a few other collection iterators. + /// This adaptor can only be used on iterators that implement [`PeekingNext`] + /// like [`Peekable`](core::iter::Peekable), [`PutBack`] and a few other collection iterators. /// /// The last and rejected element (first `false`) is still available when /// `peeking_take_while` is done. @@ -1557,10 +1560,20 @@ pub trait Itertools: Iterator { peeking_take_while::peeking_take_while(self, accept) } - /// Return an iterator adaptor that borrows from a `Clone`-able iterator + /// Return an iterator adaptor thar borrows from this iterator and takes items while the closure returns [`Either::Left`]. + /// If it returns [`Either::Right`] then that value is put back into the iterator (which is implementation dependent) and the iterator returns none + fn peeking_map_while(&mut self, accept: F) -> PeekingMapWhile<'_, Self, F> + where + Self: Sized + PeekingMapNext, + F: FnMut(Self::Item) -> Either, + { + peeking_map_while::PeekingMapWhile { iter: self, accept } + } + + /// Return an iterator adaptor that borrows from a [`Clone`]-able iterator /// to only pick off elements while the predicate `accept` returns `true`. /// - /// It uses the `Clone` trait to restore the original iterator so that the + /// It uses the [`Clone`] trait to restore the original iterator so that the /// last and rejected element (first `false`) is still available when /// `take_while_ref` is done. /// diff --git a/src/peeking_map_while.rs b/src/peeking_map_while.rs new file mode 100644 index 000000000..6f80f47fc --- /dev/null +++ b/src/peeking_map_while.rs @@ -0,0 +1,86 @@ +use core::{iter::Peekable, mem}; + +use either::Either; + +use crate::PutBack; + +/// An iterator adaptor that takes items while a closure returns [`Either::Left`]. +/// +/// See [`.peeking_map_while()`](crate::Itertools::peeking_map_while) +/// for more information. +#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] +pub struct PeekingMapWhile<'a, I, F> { + pub(crate) iter: &'a mut I, + pub(crate) accept: F, +} + +impl Iterator for PeekingMapWhile<'_, I, F> +where + I: PeekingMapNext, + F: FnMut(I::Item) -> Either, +{ + type Item = O; + + fn next(&mut self) -> Option { + self.iter.peeking_map_next(&mut self.accept) + } +} + +/// A trait used in [`PeekingMapWhile`] +pub trait PeekingMapNext: Iterator { + /// If it returns [`Either::Right`] then it should be reinserted to the iterator otherwise it should be returned + fn peeking_map_next( + &mut self, + accept: impl FnMut(Self::Item) -> Either, + ) -> Option; +} + +impl PeekingMapNext for Peekable +where + I: Iterator, + I::Item: Default, +{ + fn peeking_map_next( + &mut self, + mut accept: impl FnMut(I::Item) -> Either, + ) -> Option { + let dest = self.peek_mut()?; + let item = mem::take(dest); + let out = accept(item); + match out { + Either::Left(out) => { + self.next(); + Some(out) + } + Either::Right(item) => { + *dest = item; + None + } + } + } +} + +impl PeekingMapNext for PutBack +where + I: Iterator, +{ + fn peeking_map_next( + &mut self, + mut accept: impl FnMut(I::Item) -> Either, + ) -> Option { + match accept(self.next()?) { + Either::Left(out) => Some(out), + Either::Right(item) => { + self.put_back(item); + None + } + } + } +} + +impl<'a, I, F> std::fmt::Debug for PeekingMapWhile<'a, I, F> +where + I: Iterator + std::fmt::Debug + 'a, +{ + debug_fmt_fields!(PeekingTakeWhile, iter); +} diff --git a/tests/peeking_map_while.rs b/tests/peeking_map_while.rs new file mode 100644 index 000000000..067f43881 --- /dev/null +++ b/tests/peeking_map_while.rs @@ -0,0 +1,28 @@ +use either::Either; +use itertools::{put_back, Itertools}; + +#[test] +fn peeking_map_while_peekable() { + let mut r = (0..10).peekable(); + let last = r + .peeking_map_while(|x| match x { + 0..3 => Either::Left(x * 2), + _ => Either::Right(x), + }) + .last(); + assert_eq!(last, Some(4)); + assert_eq!(r.next(), Some(3)); +} + +#[test] +fn peeking_map_while_put_back() { + let mut r = put_back(0..10); + let last = r + .peeking_map_while(|x| match x { + 0..3 => Either::Left(x * 2), + _ => Either::Right(x), + }) + .last(); + assert_eq!(last, Some(4)); + assert_eq!(r.next(), Some(3)); +}