tor-basic-utils 0.5.0

//! Iterator helpers for Arti.

/// Iterator extension trait to implement a counting filter.
pub trait IteratorExt: Iterator {
    /// Return an iterator that contains every member of this iterator, and
    /// which records its progress in `count`.
    ///
    /// The values in `count` are initially set to zero.  Then, every time the
    /// filter considers an item, it will either increment `count.n_accepted` or
    /// `count.n_rejected`.
    ///
    /// Note that if the iterator is dropped before it is exhausted, the count will not
    /// be complete.
    ///
    /// # Examples
    ///
    /// ```
    /// use tor_basic_utils::iter::{IteratorExt, FilterCount};
    ///
    /// let mut count = FilterCount::default();
    /// let emoji : String = "Hello 🙂 World 🌏!"
    ///     .chars()
    ///     .filter_cnt(&mut count, |ch| !ch.is_ascii())
    ///     .collect();
    /// assert_eq!(emoji, "🙂🌏");
    /// assert_eq!(count, FilterCount { n_accepted: 2, n_rejected: 14});
    /// ```
    //
    // In Arti, we mostly use this iterator for reporting issues when we're
    // unable to find a suitable relay for some purpose: it makes it easy to
    // tabulate which filters in a chain of filters rejected how many of the
    // potential candidates.
    fn filter_cnt<P>(self, count: &mut FilterCount, pred: P) -> CountingFilter<'_, P, Self>
    where
        Self: Sized,
        P: FnMut(&Self::Item) -> bool,
    {
        *count = FilterCount::default();
        CountingFilter {
            inner: self,
            pred,
            count,
        }
    }
}

impl<I> IteratorExt for I where I: Iterator {}

/// A record of how many items a [`CountingFilter`] returned by
/// [`IteratorExt::filter_cnt`] accepted and rejected.
///
/// In `tor-guardmgr` we use this type to keep track of which filters reject which guards.
//
// SEMVER NOTE: This type has public members, is exhaustive, and is re-exposed
// from various error types elsewhere in arti.  Probably you should not change
// its members.  If you do, you will need to mark it as a breaking change
// everywhere that it is re-exported.
#[derive(Copy, Clone, Default, Debug, Eq, PartialEq)]
#[allow(clippy::exhaustive_structs)]
pub struct FilterCount {
    /// The number of items that the filter considered and accepted.
    pub n_accepted: usize,
    /// The number of items that the filter considered and accepted.
    pub n_rejected: usize,
}

/// An iterator to implement [`IteratorExt::filter_cnt`].
pub struct CountingFilter<'a, P, I> {
    /// The inner iterator that we're taking items from.
    inner: I,
    /// The predicate we're using to decide which items are accepted.
    pred: P,
    /// The count of the number of items accepted and rejected so far.
    count: &'a mut FilterCount,
}

impl<'a, P, I> Iterator for CountingFilter<'a, P, I>
where
    P: FnMut(&I::Item) -> bool,
    I: Iterator,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        for item in &mut self.inner {
            if (self.pred)(&item) {
                self.count.n_accepted += 1;
                return Some(item);
            } else {
                self.count.n_rejected += 1;
            }
        }
        None
    }
}

impl FilterCount {
    /// Return a wrapper that can be displayed as the fraction of rejected items.
    ///
    /// # Example
    ///
    /// ```
    /// # use tor_basic_utils::iter::{IteratorExt, FilterCount};
    /// let mut count = FilterCount::default();
    /// let sum_of_evens : u32 = (1..=10)
    ///     .filter_cnt(&mut count, |x| *x % 2 == 0)
    ///     .sum();
    /// assert_eq!(format!("Rejected {} as odd", count.display_frac_rejected()),
    ///     "Rejected 5/10 as odd".to_string());
    /// ```
    pub fn display_frac_rejected(&self) -> DisplayFracRejected<'_> {
        DisplayFracRejected(self)
    }

    /// Count and return the provided boolean value.
    ///
    /// This is an alternative way to use `FilterCount` when you have to provide
    /// a function that takes a predicate rather than a member of an iterator
    /// chain.
    ///
    /// # Example
    ///
    /// ```
    /// # use tor_basic_utils::iter::FilterCount;
    /// let mut count = FilterCount::default();
    /// let mut emoji = "Hello 🙂 World 🌏!".to_string();
    /// emoji.retain(|ch| count.count(!ch.is_ascii()));
    /// assert_eq!(emoji, "🙂🌏");
    /// assert_eq!(count, FilterCount { n_accepted: 2, n_rejected: 14});
    /// ```
    pub fn count(&mut self, accept: bool) -> bool {
        if accept {
            self.n_accepted += 1;
        } else {
            self.n_rejected += 1;
        }
        accept
    }
}

/// Return value from [`FilterCount::display_frac_rejected`].
#[derive(Debug, Clone)]
pub struct DisplayFracRejected<'a>(&'a FilterCount);

impl<'a> std::fmt::Display for DisplayFracRejected<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}/{}",
            self.0.n_rejected,
            self.0.n_accepted + self.0.n_rejected
        )
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod test {
    use super::*;

    #[test]
    fn counting_filter() {
        let mut count = FilterCount::default();
        let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9];
        let first_even = v
            .iter()
            .filter_cnt(&mut count, |val| **val % 2 == 0)
            .next()
            .unwrap();
        assert_eq!(*first_even, 2);
        assert_eq!(count.n_accepted, 1);
        assert_eq!(count.n_rejected, 1);

        let sum_even: usize = v.iter().filter_cnt(&mut count, |val| **val % 2 == 0).sum();
        assert_eq!(sum_even, 20);
        assert_eq!(count.n_accepted, 4);
        assert_eq!(count.n_rejected, 5);
    }

    #[test]
    fn counting_with_predicates() {
        let mut count = FilterCount::default();
        let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9];
        let first_even = v.iter().find(|val| count.count(**val % 2 == 0)).unwrap();
        assert_eq!(*first_even, 2);
        assert_eq!(count.n_accepted, 1);
        assert_eq!(count.n_rejected, 1);

        let mut count = FilterCount::default();
        let sum_even: usize = v.iter().filter(|val| count.count(**val % 2 == 0)).sum();
        assert_eq!(sum_even, 20);
        assert_eq!(count.n_accepted, 4);
        assert_eq!(count.n_rejected, 5);
    }

    #[test]
    fn fooz() {}
}