Iterator extension AllTheSame

[Schard]

Today I had a use case, where I needed to check if all elements in an iterator are the same.
I looked at the standard library and itertools but did not find anything that resembles my use case.
Here's the extension trait that solves my issue:

use std::collections::HashSet;
use std::hash::Hash;

/// An iterator extension that checks if all elements are the same.
///
/// # Examples
/// ```
/// use all_the_same::AllTheSame;
///
/// let v = vec![1, 1, 1, 1];
/// assert_eq!(v.into_iter().all_the_same(), Ok(1));
///
/// let v = vec![1, 2, 1, 1];
/// assert_eq!(v.into_iter().all_the_same(), Err([1, 2].iter().cloned().collect()));
pub trait AllTheSame: Iterator {
    /// Returns the first element if all elements are the same.
    ///
    /// # Errors
    /// Returns a set of all different elements if they are not the same.
    /// Returns an empty set if the iterator is empty.
    fn all_the_same(self) -> Result<Self::Item, HashSet<Self::Item>>;
}

impl<T> AllTheSame for T
where
    T: Iterator,
    <T as Iterator>::Item: Eq + Hash,
{
    fn all_the_same(mut self) -> Result<Self::Item, HashSet<Self::Item>> {
        let mut seen = HashSet::new();

        if let Some(first) = self.next() {
            for other in self {
                if other != first {
                    seen.insert(other);
                }
            }

            if seen.is_empty() {
                Ok(first)
            } else {
                seen.insert(first);
                Err(seen)
            }
        } else {
            Err(seen)
        }
    }
}

Did I reinvent the wheel? Can I improve the code?

[VoreLuka]

Maybe something along the lines of

first = iter.first().unwrap();
return iter.all(|element| element == first);

[steffahn]

Here’s an alternative way of doing the implementation that I could come up with

let first = self.next().ok_or_else(HashSet::new)?;
match self.find(|other| *other != first) {
    None => Ok(first),
    Some(first_other) => Err(HashSet::from_iter(
        [first, first_other].into_iter().chain(self),
    )),
}

I went for a more ”functional” approach here, which defines the return values in one step in each case, not build them up in multiple mutations.

Beyond a different style (which is just a matter of taste), a small advantage I could find is that you don't loop over the whole list doing for every item 2 steps, first checking other != first and then doing the HashSet::insert, but instead the hashset-lookup can handle both at once more efficiently if first was already part of the set.

Some things I like about your code already: It makes sure to have consistent behavior for choosing the representative of equal items (it seems HashSet always keeps the first value if multiple are equal, and your code also keeps the first value of all the ones equal to first). It also effectively makes sure that no hashset is allocated unless there’s at least 2 different items.

---

A general possible issue with this function of yours is that if one isn’t interested in the Err variant, the function does additional work beyond determining whether all items are the same (in which case it could return earlier). That being said, that’s mainly a naming and/or documentation issue, and usage of Result where the Err variant is considered important enough to be fed with additional “expensive”-to-compute information is completely fine, as long as the user is aware and consider the cost insignificant or are aware that they should use different API if they aren’t interested in the Err value.

---

Some more variation on my code: With a free-standing chain function, like perhaps eventually part of std and currently part of itertools, the expression of [first, first_other].into_iter().chain(self) can become a bit nicer as chain([first, first_other], self). The usage of HashSet::from_iter can also be replaced with a call to .collect(), the latter is shorter, the former more explicitly makes the type visible.

[Schard]

I don't want my code to panic, but your example made me thinking:

Rust Playground

I still don't know whether I want to pass on having the set of unique values, though.

[steffahn]

I just found Itertools::all_equal_value, which does half of my code’s functionality already. Using this, the code can become the following:

use itertools::Itertools;
use itertools::chain;

self.all_equal_value().map_err(|e| match e {
    None => HashSet::new(),
    Some((fst, fst_other)) => chain([fst, fst_other], self).collect(),
})

It’s not much shorter, because the behavior of all the 2 or 3 relevant cases still need to be defined (especially empty list becoming empty HashSet and 2 mismatches found becoming a HashSet containing these and all other unique iterator elements).

[Schard]

Darn it. I thought I looked into itertools thoroughly enough, but I must have overlooked that.
Having the first diverging value is probably a good tradeoff for my use case, rather than rolling NIH code.