Merge pull request #37 from stepancheg/unique

.unique() and .unique_by(..) operations

Author: bluss

src/adaptors.rs

@@ -11,6 +11,8 @@ use std::num::One;
 use std::ops::Add;
 use std::cmp::Ordering;
 use std::iter::{Fuse, Peekable};
+use std::collections::HashSet;
+use std::hash::Hash;
 use Itertools;
 use size_hint;
 
@@ -972,3 +974,55 @@ impl<I> Iterator for Combinations<I> where I: Iterator + Clone, I::Item: Clone{
         size_hint::add((lo / 2, hi.map(|hi| hi / 2)), extra)
     }
 }
+
+/// An iterator adapter to filter non-unique elements.
+///
+/// See [*.unique()*](trait.Itertools.html#method.unique) for more information.
+#[derive(Clone)]
+pub struct UniqueBy<I: Iterator, V, F> {
+    iter: I,
+    used: HashSet<V>,
+    f: F,
+}
+
+impl<I: Iterator, V, F> UniqueBy<I, V, F> where V: Clone + Eq + Hash, F: FnMut(&I::Item) -> V {
+    /// Create a new **UniqueBy** iterator.
+    pub fn new(iter: I, f: F) -> UniqueBy<I, V, F> {
+        UniqueBy {
+            iter: iter,
+            used: HashSet::new(),
+            f: f,
+        }
+    }
+}
+
+impl<I, V, F> Iterator for UniqueBy<I, V, F> where
+    I: Iterator,
+    V: Clone + Eq + Hash,
+    F: FnMut(&I::Item) -> V
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<I::Item> {
+        loop {
+            match self.iter.next() {
+                None => return None,
+                Some(v) => {
+                    let key = (self.f)(&v);
+                    if self.used.insert(key) {
+                        return Some(v);
+                    }
+                }
+            }
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (0, self.iter.size_hint().1)
+    }
+}
+
+/// An iterator adapter to filter non-unique elements.
+pub type Unique<I> where I: Iterator =
+    UniqueBy<I, I::Item, fn(&I::Item) -> I::Item>;

src/lib.rs

@@ -36,6 +36,7 @@ use std::iter::{self, IntoIterator};
 use std::fmt::Write;
 use std::cmp::Ordering;
 use std::fmt;
+use std::hash::Hash;
 
 pub use adaptors::{
     Interleave,
@@ -54,6 +55,8 @@ pub use adaptors::{
     Coalesce,
     CoalesceFn,
     Combinations,
+    Unique,
+    UniqueBy,
 };
 #[cfg(feature = "unstable")]
 pub use adaptors::EnumerateFrom;
@@ -625,6 +628,45 @@ pub trait Itertools : Iterator {
         Coalesce::new(self, eq)
     }
 
+    /// Filter non-unique elements from the iterator.
+    ///
+    /// Copies of visited elements are stored in a hash set in the
+    /// iterator.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let data = vec![10, 20, 30, 20, 40, 10, 50];
+    /// itertools::assert_equal(data.into_iter().unique(),
+    ///                         vec![10, 20, 30, 40, 50]);
+    /// ```
+    fn unique(self) -> Unique<Self> where
+        Self: Sized,
+        Self::Item: Clone + Eq + Hash,
+    {
+        self.unique_by(Clone::clone)
+    }
+
+    /// Filter non-unique elements from the iterator.
+    ///
+    /// Elemens are considered the same if supplied function returns
+    /// equal values for them. Those values are stored in a hash set in
+    /// the iterator.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let data = vec!["a", "bb", "aa", "c", "ccc"];
+    /// itertools::assert_equal(data.into_iter().unique_by(|s| s.len()),
+    ///                         vec!["a", "bb", "ccc"]);
+    /// ```
+    fn unique_by<V, F>(self, f: F) -> UniqueBy<Self, V, F> where
+        Self: Sized,
+        V: Clone + Eq + Hash,
+        F: FnMut(&Self::Item) -> V
+    {
+        UniqueBy::new(self, f)
+    }
 
     /// Return an iterator adaptor that joins together adjacent slices if possible.
     ///

tests/tests.rs

@@ -202,6 +202,23 @@ fn dedup() {
     it::assert_equal(ys.iter(), xs.iter().dedup());
 }
 
+#[test]
+fn unique_by() {
+    let xs = ["aaa", "bbbbb", "aa", "ccc", "bbbb", "aaaaa", "cccc"];
+    let ys = ["aaa", "bbbbb", "ccc"];
+    it::assert_equal(ys.iter(), xs.iter().unique_by(|x| x[..2].to_string()));
+}
+
+#[test]
+fn unique() {
+    let xs = [0, 1, 2, 3, 2, 1, 3];
+    let ys = [0, 1, 2, 3];
+    it::assert_equal(ys.iter(), xs.iter().unique());
+    let xs = [0, 1];
+    let ys = [0, 1];
+    it::assert_equal(ys.iter(), xs.iter().unique());
+}
+
 #[test]
 fn batching() {
     let xs = [0, 1, 2, 1, 3];