diff --git a/src/lib.rs b/src/lib.rs index 0c6af37e4..407a58cdc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -130,6 +130,7 @@ pub mod structs { pub use crate::peek_nth::PeekNth; pub use crate::pad_tail::PadUsing; pub use crate::peeking_take_while::PeekingTakeWhile; + pub use crate::peeking_fold_while::PeekingFoldWhile; #[cfg(feature = "use_alloc")] pub use crate::permutations::Permutations; pub use crate::process_results_impl::ProcessResults; @@ -210,6 +211,7 @@ mod multipeek_impl; mod pad_tail; #[cfg(feature = "use_alloc")] mod peek_nth; +mod peeking_fold_while; mod peeking_take_while; #[cfg(feature = "use_alloc")] mod permutations; @@ -1214,6 +1216,39 @@ pub trait Itertools : Iterator { peeking_take_while::peeking_take_while(self, accept) } + /// An iterator method that applies a function to each element + /// as long as it returns successfully, producing a single value. + /// + /// Unlike `try_fold()`, `peeking_fold_while()` does not consume the element + /// that causes the function to short-circuit. + /// + /// `peeking_fold_while()` is particularly useful when the short-circuit + /// condition depends on the accumulated value. + /// + /// # Example + /// ``` + /// let a = [10, 20, 30, 100, 40, 50]; + /// + /// // Using `try_fold()` + /// let mut it = a.iter(); + /// let sum = it.try_fold(0i8, |acc, &x| acc.checked_add(x).ok_or(acc)); + /// assert_eq!(sum, Err(60)); + /// assert_eq!(it.next(), Some(&40)); + /// + /// // Using `peeking_fold_while()` + /// use itertools::Itertools; + /// let mut it = a.iter().peekable(); + /// let sum = it.peeking_fold_while(0i8, |acc, &&x| acc.checked_add(x).ok_or(acc)); + /// assert_eq!(sum, Err(60)); + /// assert_eq!(it.next(), Some(&100)); + /// ``` + fn peeking_fold_while(&mut self, init: T, f: F) -> Result + where Self: PeekingFoldWhile, + F: FnMut(T, &Self::Item) -> Result, + { + PeekingFoldWhile::peeking_fold_while(self, init, f) + } + /// Return an iterator adaptor that borrows from a `Clone`-able iterator /// to only pick off elements while the predicate `accept` returns `true`. /// diff --git a/src/peeking_fold_while.rs b/src/peeking_fold_while.rs new file mode 100644 index 000000000..4da863d30 --- /dev/null +++ b/src/peeking_fold_while.rs @@ -0,0 +1,71 @@ +use std::iter::Peekable; +use crate::PutBack; +#[cfg(feature = "use_std")] +use crate::PutBackN; + +/// A trait for folding an iterator by peeking at its elements before +/// consuming them. +pub trait PeekingFoldWhile : Iterator { + /// An iterator method that applies a function to each element + /// as long as it returns successfully, producing a single value. + /// + /// See [`.peeking_fold_while()`](../trait.Itertools.html#method.peeking_fold_while) for + /// more information. + fn peeking_fold_while(&mut self, init: T, f: F) -> Result + where F: FnMut(T, &Self::Item) -> Result; +} + +impl PeekingFoldWhile for Peekable + where I:Iterator +{ + fn peeking_fold_while(&mut self, init: T, mut f: F) -> Result + where F: FnMut(T, &I::Item) -> Result, + { + let mut acc = init; + while let Some(x) = self.peek() { + let result = f(acc, x); + if result.is_ok() { + self.next(); + } + acc = result?; + } + Ok(acc) + } +} + +impl PeekingFoldWhile for PutBack + where I: Iterator +{ + fn peeking_fold_while(&mut self, init: T, mut f: F) -> Result + where F: FnMut(T, &I::Item) -> Result, + { + let mut acc = init; + while let Some(x) = self.next() { + let result = f(acc, &x); + if result.is_err() { + self.put_back(x); + } + acc = result?; + } + Ok(acc) + } +} + +#[cfg(feature = "use_std")] +impl PeekingFoldWhile for PutBackN + where I: Iterator +{ + fn peeking_fold_while(&mut self, init: T, mut f: F) -> Result + where F: FnMut(T, &I::Item) -> Result, + { + let mut acc = init; + while let Some(x) = self.next() { + let result = f(acc, &x); + if result.is_err() { + self.put_back(x); + } + acc = result?; + } + Ok(acc) + } +} diff --git a/tests/peeking_fold_while.rs b/tests/peeking_fold_while.rs new file mode 100644 index 000000000..fc7621f76 --- /dev/null +++ b/tests/peeking_fold_while.rs @@ -0,0 +1,58 @@ +extern crate itertools; + +use itertools::Itertools; +use itertools::{put_back, put_back_n}; + +#[test] +fn peeking_fold_while_peekable_consumes_all() { + let a = [10, 20, 30]; + let mut it = a.iter().peekable(); + let sum = it.peeking_fold_while(0i8, |acc, &&x| acc.checked_add(x).ok_or(acc)); + assert_eq!(sum, Ok(60)); + assert_eq!(it.next(), None); +} + +#[test] +fn peeking_fold_while_peekable_consumes_some() { + let a = [10, 20, 30, 100, 40, 50]; + let mut it = a.iter().peekable(); + let sum = it.peeking_fold_while(0i8, |acc, &&x| acc.checked_add(x).ok_or(acc)); + assert_eq!(sum, Err(60)); + assert_eq!(it.next(), Some(&100)); +} + +#[test] +fn peeking_fold_while_put_back_consumes_all() { + let a = [10, 20, 30]; + let mut it = put_back(a.iter()); + let sum = it.peeking_fold_while(0i8, |acc, &&x| acc.checked_add(x).ok_or(acc)); + assert_eq!(sum, Ok(60)); + assert_eq!(it.next(), None); +} + +#[test] +fn peeking_fold_while_put_back_consumes_some() { + let a = [10, 20, 30, 100, 40, 50]; + let mut it = put_back(a.iter()); + let sum = it.peeking_fold_while(0i8, |acc, &&x| acc.checked_add(x).ok_or(acc)); + assert_eq!(sum, Err(60)); + assert_eq!(it.next(), Some(&100)); +} + +#[test] +fn peeking_fold_while_put_back_n_consumes_all() { + let a = [10, 20, 30]; + let mut it = put_back_n(a.iter()); + let sum = it.peeking_fold_while(0i8, |acc, &&x| acc.checked_add(x).ok_or(acc)); + assert_eq!(sum, Ok(60)); + assert_eq!(it.next(), None); +} + +#[test] +fn peeking_fold_while_put_back_n_consumes_some() { + let a = [10, 20, 30, 100, 40, 50]; + let mut it = put_back(a.iter()); + let sum = it.peeking_fold_while(0i8, |acc, &&x| acc.checked_add(x).ok_or(acc)); + assert_eq!(sum, Err(60)); + assert_eq!(it.next(), Some(&100)); +}