itertools/
unique_impl.rs

1
2use std::collections::HashMap;
3use std::collections::hash_map::{Entry};
4use std::hash::Hash;
5use std::fmt;
6
7/// An iterator adapter to filter out duplicate elements.
8///
9/// See [`.unique_by()`](../trait.Itertools.html#method.unique) for more information.
10#[derive(Clone)]
11#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
12pub struct UniqueBy<I: Iterator, V, F> {
13    iter: I,
14    // Use a hashmap for the entry API
15    used: HashMap<V, ()>,
16    f: F,
17}
18
19impl<I, V, F> fmt::Debug for UniqueBy<I, V, F>
20    where I: Iterator + fmt::Debug,
21          V: fmt::Debug + Hash + Eq,
22{
23    debug_fmt_fields!(UniqueBy, iter, used);
24}
25
26/// Create a new `UniqueBy` iterator.
27pub fn unique_by<I, V, F>(iter: I, f: F) -> UniqueBy<I, V, F>
28    where V: Eq + Hash,
29          F: FnMut(&I::Item) -> V,
30          I: Iterator,
31{
32    UniqueBy {
33        iter: iter,
34        used: HashMap::new(),
35        f: f,
36    }
37}
38
39// count the number of new unique keys in iterable (`used` is the set already seen)
40fn count_new_keys<I, K>(mut used: HashMap<K, ()>, iterable: I) -> usize
41    where I: IntoIterator<Item=K>,
42          K: Hash + Eq,
43{
44    let iter = iterable.into_iter();
45    let current_used = used.len();
46    used.extend(iter.map(|key| (key, ())));
47    used.len() - current_used
48}
49
50impl<I, V, F> Iterator for UniqueBy<I, V, F>
51    where I: Iterator,
52          V: Eq + Hash,
53          F: FnMut(&I::Item) -> V
54{
55    type Item = I::Item;
56
57    fn next(&mut self) -> Option<I::Item> {
58        while let Some(v) = self.iter.next() {
59            let key = (self.f)(&v);
60            if self.used.insert(key, ()).is_none() {
61                return Some(v);
62            }
63        }
64        None
65    }
66
67    #[inline]
68    fn size_hint(&self) -> (usize, Option<usize>) {
69        let (low, hi) = self.iter.size_hint();
70        ((low > 0 && self.used.is_empty()) as usize, hi)
71    }
72
73    fn count(self) -> usize {
74        let mut key_f = self.f;
75        count_new_keys(self.used, self.iter.map(move |elt| key_f(&elt)))
76    }
77}
78
79impl<I> Iterator for Unique<I>
80    where I: Iterator,
81          I::Item: Eq + Hash + Clone
82{
83    type Item = I::Item;
84
85    fn next(&mut self) -> Option<I::Item> {
86        while let Some(v) = self.iter.iter.next() {
87            if let Entry::Vacant(entry) = self.iter.used.entry(v) {
88                let elt = entry.key().clone();
89                entry.insert(());
90                return Some(elt);
91            }
92        }
93        None
94    }
95
96    #[inline]
97    fn size_hint(&self) -> (usize, Option<usize>) {
98        let (low, hi) = self.iter.iter.size_hint();
99        ((low > 0 && self.iter.used.is_empty()) as usize, hi)
100    }
101
102    fn count(self) -> usize {
103        count_new_keys(self.iter.used, self.iter.iter)
104    }
105}
106
107/// An iterator adapter to filter out duplicate elements.
108///
109/// See [`.unique()`](../trait.Itertools.html#method.unique) for more information.
110#[derive(Clone)]
111#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
112pub struct Unique<I: Iterator> {
113    iter: UniqueBy<I, I::Item, ()>,
114}
115
116impl<I> fmt::Debug for Unique<I>
117    where I: Iterator + fmt::Debug,
118          I::Item: Hash + Eq + fmt::Debug,
119{
120    debug_fmt_fields!(Unique, iter);
121}
122
123pub fn unique<I>(iter: I) -> Unique<I>
124    where I: Iterator,
125          I::Item: Eq + Hash,
126{
127    Unique {
128        iter: UniqueBy {
129            iter: iter,
130            used: HashMap::new(),
131            f: (),
132        }
133    }
134}