Implement peeking_fold_while

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<T, E, F>(&mut self, init: T, f: F) -> Result<T, E>
+        where Self: PeekingFoldWhile,
+              F: FnMut(T, &Self::Item) -> Result<T, E>,
+    {
+        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`.
     ///

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<T, E, F>(&mut self, init: T, f: F) -> Result<T, E>
+        where F: FnMut(T, &Self::Item) -> Result<T, E>;
+}
+
+impl<I> PeekingFoldWhile for Peekable<I>
+    where I:Iterator
+{
+    fn peeking_fold_while<T, E, F>(&mut self, init: T, mut f: F) -> Result<T, E>
+        where F: FnMut(T, &I::Item) -> Result<T, E>,
+    {
+        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<I> PeekingFoldWhile for PutBack<I>
+    where I: Iterator
+{
+    fn peeking_fold_while<T, E, F>(&mut self, init: T, mut f: F) -> Result<T, E>
+        where F: FnMut(T, &I::Item) -> Result<T, E>,
+    {
+        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<I> PeekingFoldWhile for PutBackN<I>
+    where I: Iterator
+{
+    fn peeking_fold_while<T, E, F>(&mut self, init: T, mut f: F) -> Result<T, E>
+        where F: FnMut(T, &I::Item) -> Result<T, E>,
+    {
+        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)
+    }
+}

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));
+}