itertools/
unique_impl.rs

1use std::collections::HashMap;
2use std::collections::hash_map::Entry;
3use std::hash::Hash;
4use std::fmt;
5use std::iter::FusedIterator;
6
7/// An iterator adapter to filter out duplicate elements.
8///
9/// See [`.unique_by()`](crate::Itertools::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 in order to prevent hashing twice.
15    // This can maybe be replaced with a HashSet once `get_or_insert_with`
16    // or a proper Entry API for Hashset is stable and meets this msrv
17    used: HashMap<V, ()>,
18    f: F,
19}
20
21impl<I, V, F> fmt::Debug for UniqueBy<I, V, F>
22    where I: Iterator + fmt::Debug,
23          V: fmt::Debug + Hash + Eq,
24{
25    debug_fmt_fields!(UniqueBy, iter, used);
26}
27
28/// Create a new `UniqueBy` iterator.
29pub fn unique_by<I, V, F>(iter: I, f: F) -> UniqueBy<I, V, F>
30    where V: Eq + Hash,
31          F: FnMut(&I::Item) -> V,
32          I: Iterator,
33{
34    UniqueBy {
35        iter,
36        used: HashMap::new(),
37        f,
38    }
39}
40
41// count the number of new unique keys in iterable (`used` is the set already seen)
42fn count_new_keys<I, K>(mut used: HashMap<K, ()>, iterable: I) -> usize
43    where I: IntoIterator<Item=K>,
44          K: Hash + Eq,
45{
46    let iter = iterable.into_iter();
47    let current_used = used.len();
48    used.extend(iter.map(|key| (key, ())));
49    used.len() - current_used
50}
51
52impl<I, V, F> Iterator for UniqueBy<I, V, F>
53    where I: Iterator,
54          V: Eq + Hash,
55          F: FnMut(&I::Item) -> V
56{
57    type Item = I::Item;
58
59    fn next(&mut self) -> Option<Self::Item> {
60        while let Some(v) = self.iter.next() {
61            let key = (self.f)(&v);
62            if self.used.insert(key, ()).is_none() {
63                return Some(v);
64            }
65        }
66        None
67    }
68
69    #[inline]
70    fn size_hint(&self) -> (usize, Option<usize>) {
71        let (low, hi) = self.iter.size_hint();
72        ((low > 0 && self.used.is_empty()) as usize, hi)
73    }
74
75    fn count(self) -> usize {
76        let mut key_f = self.f;
77        count_new_keys(self.used, self.iter.map(move |elt| key_f(&elt)))
78    }
79}
80
81impl<I, V, F> DoubleEndedIterator for UniqueBy<I, V, F>
82    where I: DoubleEndedIterator,
83          V: Eq + Hash,
84          F: FnMut(&I::Item) -> V
85{
86    fn next_back(&mut self) -> Option<Self::Item> {
87        while let Some(v) = self.iter.next_back() {
88            let key = (self.f)(&v);
89            if self.used.insert(key, ()).is_none() {
90                return Some(v);
91            }
92        }
93        None
94    }
95}
96
97impl<I, V, F> FusedIterator for UniqueBy<I, V, F>
98    where I: FusedIterator,
99          V: Eq + Hash,
100          F: FnMut(&I::Item) -> V
101{}
102
103impl<I> Iterator for Unique<I>
104    where I: Iterator,
105          I::Item: Eq + Hash + Clone
106{
107    type Item = I::Item;
108
109    fn next(&mut self) -> Option<Self::Item> {
110        while let Some(v) = self.iter.iter.next() {
111            if let Entry::Vacant(entry) = self.iter.used.entry(v) {
112                let elt = entry.key().clone();
113                entry.insert(());
114                return Some(elt);
115            }
116        }
117        None
118    }
119
120    #[inline]
121    fn size_hint(&self) -> (usize, Option<usize>) {
122        let (low, hi) = self.iter.iter.size_hint();
123        ((low > 0 && self.iter.used.is_empty()) as usize, hi)
124    }
125
126    fn count(self) -> usize {
127        count_new_keys(self.iter.used, self.iter.iter)
128    }
129}
130
131impl<I> DoubleEndedIterator for Unique<I>
132    where I: DoubleEndedIterator,
133          I::Item: Eq + Hash + Clone
134{
135    fn next_back(&mut self) -> Option<Self::Item> {
136        while let Some(v) = self.iter.iter.next_back() {
137            if let Entry::Vacant(entry) = self.iter.used.entry(v) {
138                let elt = entry.key().clone();
139                entry.insert(());
140                return Some(elt);
141            }
142        }
143        None
144    }
145}
146
147impl<I> FusedIterator for Unique<I>
148    where I: FusedIterator,
149          I::Item: Eq + Hash + Clone
150{}
151
152/// An iterator adapter to filter out duplicate elements.
153///
154/// See [`.unique()`](crate::Itertools::unique) for more information.
155#[derive(Clone)]
156#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
157pub struct Unique<I: Iterator> {
158    iter: UniqueBy<I, I::Item, ()>,
159}
160
161impl<I> fmt::Debug for Unique<I>
162    where I: Iterator + fmt::Debug,
163          I::Item: Hash + Eq + fmt::Debug,
164{
165    debug_fmt_fields!(Unique, iter);
166}
167
168pub fn unique<I>(iter: I) -> Unique<I>
169    where I: Iterator,
170          I::Item: Eq + Hash,
171{
172    Unique {
173        iter: UniqueBy {
174            iter,
175            used: HashMap::new(),
176            f: (),
177        }
178    }
179}