itertools/
unique_impl.rs

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