for_each_count 0.1.0

//! Extremely simple crate providing an extension trait for the standard [Iterator] trait
//! with two utility methods:
//!
//! - [`for_each_count`]: Similar to [`for_each()`] but returns the count of processed elements
//! - [`try_for_each_count`]: Similar to [`try_for_each()`] but returns the count of processed elements in the [`Ok`] case.
//!
//! Both methods solve the *common*[^2] use case of applying a transformation to iterator elements and then returning the length of the iterator.
//! While this is trivial to achieve using the standard methods[^1], I find these extension methods to be more ergonomic, and they are available in the big extension crates, like `Itertools`.
//!
//! The provided extension methods try to mimic the way they'd be _likely_ implemented in core: based on [`fold()`] and [`try_fold()`].
//! This seems to be the preferred strategy in the standard library. [`for_each()`], [`try_for_each()`] and [`count()`] are all implemented like this at this time.
//!
//! [^1]: Like using [Iterator::inspect], or [Iterator::map] and [Iterator::count], or keeping a mutable count outside the closure passed to `for_each`.
//! [^2]: Maybe not *that* common, but I've needed it a couple of times. [Other people too].
//!
//! [Other people too]: <https://stackoverflow.com/questions/73842276/how-to-for-each-and-then-count-them-in-rust>
//!
//! [`for_each_count`]: IteratorForEachCount::for_each_count
//! [`try_for_each_count`]: IteratorForEachCount::try_for_each_count
//! [`for_each()`]: Iterator::for_each
//! [`try_for_each()`]: Iterator::try_for_each
//! [`count()`]: Iterator::count
//! [`fold()`]: Iterator::fold
//! [`try_fold()`]: Iterator::try_fold
//!
//! # Examples
//!
//! ```
//! use for_each_count::IteratorForEachCount;
//!
//! let mut sum = 0;
//! let count = (1..=5).for_each_count(|x| sum += x);
//! assert_eq!(count, 5);
//! assert_eq!(sum, 15);
//! ```

pub trait IteratorForEachCount: Iterator {
    /// Calls a closure on each element of an iterator and returns the number
    /// of times that the closure has been called.
    ///
    /// This is equivalent to the standard [`for_each`](Iterator::for_each) iterator
    /// method but also returns the number of times that the closure has been executed.
    /// That is, basically a [`for_each`](Iterator::for_each) and a [`count`] in a single iteration.
    ///
    ///
    /// [`count`]: Iterator::count
    ///
    /// # Examples
    ///
    /// Basic usage:
    ///
    /// ```
    /// use std::sync::mpsc::channel;
    /// use for_each_count::IteratorForEachCount;
    ///
    /// let (tx, rx) = channel();
    /// let count = (0..5).map(|x| x * 2 + 1)
    ///                   .for_each_count(move |x| tx.send(x).unwrap());
    ///
    /// let v: Vec<_> = rx.iter().collect();
    /// assert_eq!(v, vec![1, 3, 5, 7, 9]);
    /// assert_eq!(count, v.len());
    /// ```
    ///
    /// For such a small example, a `for` loop may be cleaner, but `for_each_count`
    /// might be preferable to keep a functional style with longer iterators:
    ///
    /// ```
    /// use for_each_count::IteratorForEachCount;
    ///
    /// let count = (0..5).flat_map(|x| x * 100 .. x * 110)
    ///                   .enumerate()
    ///                   .filter(|&(i, x)| (i + x) % 3 == 0)
    ///                   .for_each_count(|(i, x)| println!("{i}:{x}"));
    ///
    /// assert_eq!(count, 33);
    /// ```
    fn for_each_count<F>(self, f: F) -> usize
    where
        Self: Sized,
        F: FnMut(Self::Item),
    {
        #[inline]
        fn call<T>(mut f: impl FnMut(T)) -> impl FnMut(usize, T) -> usize {
            move |count, item| {
                f(item);
                count + 1
            }
        }

        self.fold(0_usize, call(f))
    }

    /// An iterator method that applies a fallible function to each item in the
    /// iterator, stopping at the first error and returning that error.
    /// If all the invocations are successful, returns the number of
    /// times that the function has been invoked.
    ///
    /// That is, this method is equivalent to the standard [`try_for_each()`] method
    /// plus [`count()`] in a single iteration loop.
    ///
    /// Just as the standard [`try_for_each()`], this method is short-circuiting.
    /// After the first error, the rest of iterator items are not consumed.
    ///
    /// This can also be thought of as the fallible form of [`for_each_count()`].
    ///
    /// [`try_for_each()`]: Iterator::try_for_each
    /// [`count()`]: Iterator::count
    /// [`for_each_count()`]: IteratorForEachCount::for_each_count
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::rename;
    /// use std::io::{stdout, Write};
    /// use std::path::Path;
    /// use for_each_count::IteratorForEachCount;
    ///
    /// let data = ["no_tea.txt", "stale_bread.json", "torrential_rain.png"];
    ///
    /// let res = data.iter().try_for_each_count(|x| writeln!(stdout(), "{x}"));
    /// assert!(res.is_ok());
    /// assert_eq!(res.unwrap(), data.len());
    ///
    /// let mut it = data.iter().cloned();
    /// let res = it.try_for_each(|x| rename(x, Path::new(x).with_extension("old")));
    /// assert!(res.is_err());
    /// // It short-circuited, so the remaining items are still in the iterator:
    /// assert_eq!(it.next(), Some("stale_bread.json"));
    /// ```
    fn try_for_each_count<F, E>(&mut self, f: F) -> Result<usize, E>
    where
        Self: Sized,
        F: FnMut(Self::Item) -> Result<(), E>,
    {
        #[inline]
        fn call<T, E>(
            mut f: impl FnMut(T) -> Result<(), E>,
        ) -> impl FnMut(usize, T) -> Result<usize, E> {
            move |count, item| {
                f(item)?;
                Ok(count + 1)
            }
        }

        self.try_fold(0_usize, call(f))
    }
}

impl<T> IteratorForEachCount for T where T: Iterator {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_empty_iterator() {
        let count = std::iter::empty::<i32>().for_each_count(|_| {});
        assert_eq!(count, 0);
    }

    #[test]
    fn test_for_each_count_with_side_effects() {
        let mut sum = 0;
        let count = (1..=5).for_each_count(|x| sum += x);
        assert_eq!(count, 5);
        assert_eq!(sum, 15);
    }

    #[test]
    fn test_for_each_count_without_side_effects() {
        let count = (1..=3).for_each_count(|_| {});
        assert_eq!(count, 3);
    }

    #[test]
    fn test_for_each_count_with_vec() {
        let mut vec = Vec::new();
        let count = (1..=4).for_each_count(|x| vec.push(x));
        assert_eq!(count, 4);
        assert_eq!(vec, vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_for_each_count_with_filter() {
        let mut sum = 0;
        let count = (1..=10).filter(|x| x % 2 == 0).for_each_count(|x| sum += x);
        assert_eq!(count, 5);
        assert_eq!(sum, 30);
    }

    #[test]
    fn test_try_for_each_count_success() {
        let mut sum = 0;
        let count = (1..=5).try_for_each_count(|x| -> Result<(), ()> {
            sum += x;
            Ok(())
        });
        assert_eq!(count.unwrap(), 5);
        assert_eq!(sum, 15);
    }

    #[test]
    fn test_try_for_each_count_error() {
        let result = (1..=5).try_for_each_count(|x| -> Result<(), &str> {
            if x == 3 { Err("error at 3") } else { Ok(()) }
        });
        assert!(result.is_err());
        assert_eq!(result.unwrap_err(), "error at 3");
    }

    #[test]
    fn test_try_for_each_count_empty() {
        let count = std::iter::empty::<i32>().try_for_each_count(|_| -> Result<(), ()> { Ok(()) });
        assert_eq!(count.unwrap(), 0);
    }

    #[test]
    fn test_try_for_each_count_with_filter() {
        let mut sum = 0;
        let count = (1..=10)
            .filter(|x| x % 2 == 0)
            .try_for_each_count(|x| -> Result<(), ()> {
                sum += x;
                Ok(())
            });
        assert_eq!(count.unwrap(), 5);
        assert_eq!(sum, 30);
    }
}